rea-agents 6.4.0 → 6.7.0

package/bridge/android/ReaJadxBridge.java

@@ -3,17 +3,26 @@ import com.google.gson.*;
 import io.modelcontextprotocol.kotlin.sdk.server.*;
 import io.modelcontextprotocol.kotlin.sdk.types.*;
 import jadx.api.*;
+import jadx.api.plugins.input.data.ICodeReader;
+import jadx.api.plugins.input.data.IMethodRef;
+import jadx.api.plugins.input.insns.InsnIndexType;
+import jadx.api.plugins.input.insns.Opcode;
+import jadx.core.dex.info.MethodInfo;
 import jadx.core.dex.nodes.*;
 import java.io.PrintStream;
+import java.nio.charset.StandardCharsets;
 import java.util.*;
 import java.util.concurrent.CountDownLatch;
 import java.util.stream.Collectors;
 import kotlin.Unit;
+import kotlin.ResultKt;
+import kotlin.coroutines.Continuation;
+import kotlin.coroutines.CoroutineContext;
 import kotlin.coroutines.EmptyCoroutineContext;
+import kotlin.coroutines.intrinsics.IntrinsicsKt;
 import kotlinx.coroutines.BuildersKt;
 import kotlinx.io.CoreKt;
 import kotlinx.io.JvmCoreKt;
-import kotlinx.serialization.json.JsonObject;
 
 /** REA-owned metadata boundary for the unmodified jadx-headless-mcp 0.7.1 JAR. */
 class ReaJadxBridge {
@@ -22,7 +31,8 @@ class ReaJadxBridge {
   private JadxSession indexedSession;
   private Map<String, List<ClassSnapshot>> classes = Map.of();
 
-  private record MethodSnapshot(String name, String shortId, Map<String, Object> summary) {}
+  private record MethodSnapshot(String name, String rawName, String shortId, List<String> bodyAbsenceMarkers,
+      Map<String, Object> summary) {}
   private record ClassSnapshot(JavaClass node, List<MethodSnapshot> methods, Map<String, Object> summary) {}
 
   private static final class InputFailure extends IllegalArgumentException {
@@ -38,15 +48,93 @@ class ReaJadxBridge {
 
   private static CallToolResult json(Object value) { return text(JSON.toJson(value), false); }
 
+  /** Translate the typed load outcome before upstream drops its exception identity. */
+  static Object loadResult(Object result) {
+    try { ResultKt.throwOnFailure(result); }
+    catch (Throwable failure) {
+      if (failure instanceof java.util.concurrent.CancellationException
+          || failure instanceof Error && !(failure instanceof OutOfMemoryError)) return result;
+      if (failure instanceof OutOfMemoryError) {
+        String reason = failure.getMessage();
+        boolean heap = "Java heap space".equals(reason) || "GC overhead limit exceeded".equals(reason);
+        Map<String, Object> reply = new LinkedHashMap<>();
+        reply.put("kind", "jvm-memory-failure");
+        reply.put("memory_area", heap ? "heap" : "unknown");
+        reply.put("exception_type", "java.lang.OutOfMemoryError");
+        if (reason != null) reply.put("message", reason);
+        reply.put("max_heap_bytes", Runtime.getRuntime().maxMemory());
+        return text(JSON.toJson(reply), true);
+      }
+      return text("load failed: " + failure.getMessage(), true);
+    }
+    SessionHolder.LoadResult loaded = (SessionHolder.LoadResult) result;
+    Map<String, Object> reply = new LinkedHashMap<>();
+    reply.put("state", "LOADED");
+    reply.put("apk_path", loaded.getApkPath());
+    reply.put("class_count", loaded.getClassCount());
+    reply.put("raw_class_count", loaded.getRawClassCount());
+    reply.put("resource_count", loaded.getResourceCount());
+    reply.put("load_duration_ms", loaded.getLoadDurationMs());
+    reply.put("decompile_timeout_ms", loaded.getDecompileTimeoutMs());
+    reply.put("threads", loaded.getThreads());
+    reply.put("code_cache_size", loaded.getCodeCacheSize());
+    reply.put("resources", loaded.getResourceMode());
+    if (!loaded.getIncludePackages().isEmpty()) reply.put("include_packages", loaded.getIncludePackages());
+    if (!loaded.getExcludePackages().isEmpty()) reply.put("exclude_packages", loaded.getExcludePackages());
+    return json(reply);
+  }
+
+  private Object load(kotlinx.serialization.json.JsonObject arguments, SessionConfig defaults,
+      Continuation<? super CallToolResult> continuation) {
+    JsonObject args = JsonParser.parseString(arguments.toString()).getAsJsonObject();
+    String path = string(args, "path");
+    SessionConfig options = defaults.withLoadOverrides(
+        optionalString(args, "include_packages"), optionalString(args, "exclude_packages"),
+        args.has("threads") ? integer(args, "threads", 1) : null,
+        args.has("code_cache_size") ? integer(args, "code_cache_size", 64) : null,
+        optionalString(args, "resources"));
+    return loadCall(next -> holder.load(path, options, next), continuation);
+  }
+
+  static Object loadCall(kotlin.jvm.functions.Function1<Continuation<SessionHolder.LoadResult>, Object> load,
+      Continuation<? super CallToolResult> continuation) {
+    var forwarded = new java.util.concurrent.atomic.AtomicBoolean();
+    Continuation<SessionHolder.LoadResult> resumed = new Continuation<>() {
+      public CoroutineContext getContext() { return continuation.getContext(); }
+      public void resumeWith(Object result) {
+        if (!forwarded.compareAndSet(false, true)) throw new IllegalStateException("Load resumed twice");
+        Object translated;
+        try { translated = loadResult(result); }
+        catch (Throwable failure) { translated = ResultKt.createFailure(failure); }
+        // Keep the consumer outside the translation catch, including inline completion.
+        continuation.resumeWith(translated);
+      }
+    };
+    Object result;
+    try { result = load.invoke(resumed); }
+    catch (Throwable failure) {
+      if (forwarded.get()) ResultKt.throwOnFailure(ResultKt.createFailure(failure));
+      result = ResultKt.createFailure(failure);
+    }
+    if (result == IntrinsicsKt.getCOROUTINE_SUSPENDED()) return result;
+    Object translated = loadResult(result);
+    ResultKt.throwOnFailure(translated);
+    return translated;
+  }
+
+  private static String optionalString(JsonObject args, String name) {
+    return args.has(name) ? string(args, name) : null;
+  }
+
   private static String string(JsonObject args, String name) {
-    JsonElement value = JsonParser.parseString(args.toString()).getAsJsonObject().get(name);
+    JsonElement value = args.get(name);
     if (value == null || !value.isJsonPrimitive() || !value.getAsJsonPrimitive().isString())
       throw new IllegalArgumentException(name + " must be a string");
     return value.getAsString();
   }
 
   private static int integer(JsonObject args, String name, int fallback) {
-    JsonElement value = JsonParser.parseString(args.toString()).getAsJsonObject().get(name);
+    JsonElement value = args.get(name);
     if (value == null) return fallback;
     if (!value.isJsonPrimitive() || !value.getAsJsonPrimitive().isNumber())
       throw new IllegalArgumentException(name + " must be a nonnegative integer");
@@ -79,11 +167,20 @@ class ReaJadxBridge {
           ordered.sort(Comparator.comparing(MethodNode::getAlias));
           List<MethodSnapshot> methods = ordered.stream().map(method -> {
             String name = method.getAlias();
+            String shortId = method.getMethodInfo().getShortId();
+            String rawName = method.getMethodInfo().getName();
+            if (!shortId.startsWith(rawName + "("))
+              throw new IllegalStateException("JADX method ID has no raw-name prefix: " + shortId);
+            String descriptor = shortId.substring(rawName.length());
+            List<String> bodyAbsenceMarkers = new ArrayList<>();
+            if (method.getAccessFlags().isNative()) bodyAbsenceMarkers.add("native_method");
+            if (method.getAccessFlags().isAbstract()) bodyAbsenceMarkers.add("abstract_method");
             String signature = name + "(" + method.getArgTypes().stream()
                 .map(Object::toString).collect(Collectors.joining(", "))
                 + "): " + method.getReturnType().toString();
-            return new MethodSnapshot(name, method.getMethodInfo().getShortId(), Map.of(
+            return new MethodSnapshot(name, rawName, shortId, List.copyOf(bodyAbsenceMarkers), Map.of(
                 "name", name, "signature", signature,
+                "dex_descriptor", descriptor,
                 "is_constructor", method.isConstructor(), "def_pos", 0));
           }).toList();
           List<Map<String, Object>> fields = node.getFields().stream()
@@ -126,7 +223,109 @@ class ReaJadxBridge {
     return result;
   }
 
-  private CallToolResult handle(String name, JsonObject args) {
+  private JadxSession.SmartCode methodBody(JadxSession current, ClassSnapshot cls,
+      MethodSnapshot selected, JavaMethod method) {
+    // Access flags belong to the parsed member, before code generation rewrites it.
+    // These methods have no Java body; source slicing can include unrelated members.
+    if (!selected.bodyAbsenceMarkers().isEmpty())
+      return new JadxSession.SmartCode("", "java", false, selected.bodyAbsenceMarkers());
+    // Java code generation can fold all or part of <clinit> into field declarations.
+    // Its method text cannot establish coverage, even when a static block remains.
+    if (selected.shortId().equals("<clinit>()V"))
+      return exactMethodSmali(current, cls, selected, List.of("static_initializer_disassembly"));
+    String java = current.getMethodSource(method);
+    if (java.startsWith("// ERROR:")) throw new IllegalArgumentException(java);
+    // The upstream placeholder means only that Java code was omitted, including
+    // methods merged into synthetic bridges. Access flags established absence above.
+    if (java.isEmpty() || java.equals(
+        "// method exists but has no decompiled body (native/abstract)"))
+      return exactMethodSmali(current, cls, selected, List.of("java_body_unavailable"));
+    List<String> markers = DecompileFailureKt.detectDecompileFailure(java);
+    if (markers.stream().anyMatch(DecompileFailureKt.getSTRONG_FAILURE_MARKERS()::contains))
+      return exactMethodSmali(current, cls, selected, markers);
+    return new JadxSession.SmartCode(
+        TextUtilKt.truncateToBytes(java, current.getMaxSourceBytes()), "java", false, markers);
+  }
+
+  private JadxSession.SmartCode exactMethodSmali(JadxSession current, ClassSnapshot cls,
+      MethodSnapshot selected, List<String> markers) {
+    JadxSession.MethodSmali smali = current.getMethodSmali(cls.node(), selected.rawName());
+    String block = selectMethodSmali(smali, selected.shortId());
+    return new JadxSession.SmartCode(
+        TextUtilKt.truncateToBytes(block, current.getMaxSourceBytes()), "smali", true, markers);
+  }
+
+  private static String selectMethodSmali(JadxSession.MethodSmali smali, String shortId) {
+    if (!smali.getFound())
+      throw new IllegalArgumentException("JADX could not disassemble " + shortId
+          + (smali.getNote() == null ? "" : ": " + smali.getNote()));
+    // Upstream returns every same-name overload, and can return an unterminated
+    // trailing block. Match the full original identity before applying output limits.
+    List<String> matches = smali.getBlocks().stream().filter(block -> {
+      String header = block.lines().findFirst().orElse("").strip();
+      return header.startsWith(".method ") && header.endsWith(" " + shortId);
+    }).toList();
+    if (matches.size() != 1)
+      throw new IllegalArgumentException("Disassembly could not establish one selected method: " + shortId);
+    String block = matches.get(0);
+    if (!block.stripTrailing().endsWith("\n.end method"))
+      throw new IllegalArgumentException("Selected method disassembly is incomplete: " + shortId);
+    return block;
+  }
+
+  // The local transport admits 8 MiB frames. Reserve 1 KiB for the fixed
+  // JSON-RPC/CallToolResult envelope and generated numeric request id, and
+  // account for JSON text being escaped again inside TextContent.
+  private static final int MAX_XREF_TEXT_BYTES = 8 * 1024 * 1024 - 1024;
+  private static int escapedJsonBytes(Object value) {
+    return JSON.toJson(JSON.toJson(value)).getBytes(StandardCharsets.UTF_8).length;
+  }
+  private static final class InvokeBudget {
+    long remaining;
+    InvokeBudget(long remaining) { this.remaining = Math.max(0, remaining); }
+    boolean admit(Map<String, Object> site) {
+      // Check a conservative two-layer escape bound before serialization.
+      long bound = 256;
+      for (Object value : site.values())
+        if (value instanceof String text) bound += 12L * text.length();
+      if (bound > remaining) return false;
+      long bytes = escapedJsonBytes(site) - 2L + 1L;
+      if (bytes > remaining) return false;
+      remaining -= bytes;
+      return true;
+    }
+  }
+  private record DirectInvokes(List<Map<String, Object>> sites, boolean exhaustive) {}
+
+  /** Raw DEX invokes within one JADX using method; offsets are code units from its code start. */
+  private static DirectInvokes directInvokes(JavaMethod source, MethodInfo target, InvokeBudget budget) {
+    MethodNode node = source.getMethodNode();
+    ICodeReader code = node.getCodeReader();
+    if (code == null) return null;
+    List<Map<String, Object>> sites = new ArrayList<>();
+    boolean[] exhaustive = {true};
+    code.visitInstructions(insn -> {
+      insn.decode();
+      if (insn.getIndexType() != InsnIndexType.METHOD_REF) return;
+      Opcode opcode = insn.getOpcode();
+      if (!opcode.name().startsWith("INVOKE_")) return;
+      IMethodRef ref = insn.getIndexAsMethod();
+      ref.load();
+      if (!MethodInfo.fromRef(node.root(), ref).getRawFullId().equals(target.getRawFullId())) return;
+      Map<String, Object> site = Map.of(
+          "dex_entry", node.getInputFileName(),
+          "source_raw_method_id", node.getMethodInfo().getRawFullId(),
+          "dex_code_unit_offset", insn.getOffset(),
+          "dex_method_index", insn.getIndex(),
+          "opcode", opcode.name());
+      if (budget.admit(site)) sites.add(site);
+      else exhaustive[0] = false;
+    });
+    return new DirectInvokes(sites, exhaustive[0]);
+  }
+
+  private CallToolResult handle(String name, kotlinx.serialization.json.JsonObject input) {
+    JsonObject args = JsonParser.parseString(input.toString()).getAsJsonObject();
     if (name.equals("rea_jvm_status")) return json(Map.of(
         "engine_reported_version", BuildInfo.VERSION,
         "max_heap_bytes", Runtime.getRuntime().maxMemory(),
@@ -136,12 +335,25 @@ class ReaJadxBridge {
     JadxSession current = session();
     if (name.equals("list_classes")) {
       List<String> names = current.getClassFqns();
+      String query = args.has("query") ? string(args, "query") : "";
       int offset = integer(args, "offset", 0);
       int limit = integer(args, "limit", 200);
       int start = Math.min(offset, names.size());
-      int end = (int) Math.min((long) start + limit, names.size());
-      return json(Map.of("total", names.size(), "offset", offset, "limit", limit,
-          "items", names.subList(start, end)));
+      int end = start;
+      List<String> items;
+      if (query.isEmpty()) {
+        end = (int) Math.min((long) start + limit, names.size());
+        items = names.subList(start, end);
+      } else {
+        items = new ArrayList<>();
+        while (end < names.size() && items.size() < limit) {
+          String nameInInventory = names.get(end++);
+          if (nameInInventory.contains(query)) items.add(nameInInventory);
+        }
+      }
+      return json(Map.of("total_classes", names.size(), "next_offset", end,
+          "scope", "top_level", "offset", offset, "limit", limit,
+          "items", items));
     }
     ClassSnapshot cls = selected(args);
     if (name.equals("get_class_summary")) return json(cls.summary());
@@ -150,36 +362,73 @@ class ReaJadxBridge {
       int index = integer(args, "overload_index", 0);
       if (index >= overloads.size()) throw new InputFailure("overload_index", "overload_index out of range: " + index);
       JavaMethod method = method(cls, overloads.get(index));
-      JadxSession.SmartCode code = current.getMethodBodySmart(method, current.getMaxSourceBytes(), true);
+      JadxSession.SmartCode code = methodBody(current, cls, overloads.get(index), method);
       return json(Map.of("class_name", cls.summary().get("full_name"),
           "method_name", overloads.get(index).name(), "full_name", method.getFullName(),
+          "dex_descriptor", overloads.get(index).summary().get("dex_descriptor"),
           "overload_index", index, "overload_count", overloads.size(),
           "mode", code.getKind(), "fell_back", code.getFellBack(),
           "markers", code.getMarkers(), "body", code.getText()));
     }
     List<JavaNode> uses;
     String target = string(args, "class_name");
+    MethodSnapshot selectedMethod = null;
+    int overloadIndex = 0;
     if (name.equals("get_xrefs_to_class")) uses = cls.node().getUseIn();
     else {
       List<MethodSnapshot> overloads = methods(cls, args);
-      if (overloads.size() != 1) throw new InputFailure("method_name", "Method references require a unique overload");
-      target += "." + overloads.get(0).name();
-      uses = method(cls, overloads.get(0)).getUseIn();
+      if (overloads.size() != 1 && !args.has("overload_index"))
+        throw new InputFailure("overload_index", "Select one of " + overloads.size() + " method overloads");
+      overloadIndex = integer(args, "overload_index", 0);
+      if (overloadIndex >= overloads.size())
+        throw new InputFailure("overload_index", "overload_index out of range: " + overloadIndex);
+      selectedMethod = overloads.get(overloadIndex);
+      target += "." + selectedMethod.name();
+      uses = method(cls, selectedMethod).getUseIn();
     }
     // REA consumes all references without materializing source line information.
-    List<Map<String, Object>> items = uses.stream().map(node -> current.describeUsage(node, false)).toList();
-    return json(Map.of("target", target, "total", uses.size(), "count", items.size(),
-        "resolve_line", false, "items", items));
+    MethodInfo targetInfo = selectedMethod == null ? null : method(cls, selectedMethod).getMethodNode().getMethodInfo();
+    List<Map<String, Object>> items = uses.stream().map(node -> {
+      Map<String, Object> item = new LinkedHashMap<>(current.describeUsage(node, false));
+      if (targetInfo != null && node instanceof JavaMethod) {
+        item.put("direct_invokes", List.of());
+        item.put("direct_invokes_exhaustive", false);
+      }
+      return item;
+    }).toList();
+    Map<String, Object> result = new LinkedHashMap<>(Map.of("target", target,
+        "total", uses.size(), "count", items.size(), "resolve_line", false, "items", items));
+    if (selectedMethod != null) {
+      result.put("overload_index", overloadIndex);
+      result.put("reported_signature", selectedMethod.summary().get("signature"));
+    }
+    if (targetInfo != null) {
+      InvokeBudget budget = new InvokeBudget(MAX_XREF_TEXT_BYTES - escapedJsonBytes(result));
+      for (int index = 0; index < uses.size(); index++) {
+        if (!(uses.get(index) instanceof JavaMethod source)) continue;
+        DirectInvokes invokes = directInvokes(source, targetInfo, budget);
+        Map<String, Object> item = items.get(index);
+        if (invokes == null) {
+          item.remove("direct_invokes");
+          item.remove("direct_invokes_exhaustive");
+        } else {
+          item.put("direct_invokes", invokes.sites());
+          item.put("direct_invokes_exhaustive", invokes.exhaustive());
+        }
+      }
+    }
+    return json(result);
   }
 
   public static void main(String[] args) throws Exception {
     PrintStream protocolOut = System.out;
     System.setOut(System.err);
-    SessionHolder holder = new SessionHolder(new SessionConfig(
-        1048576, 0, 90000L, 1, List.of(), List.of(), 64, ResourceMode.FULL));
+    SessionConfig defaults = new SessionConfig(
+        1048576, 0, 90000L, 1, List.of(), List.of(), 64, ResourceMode.FULL);
+    SessionHolder holder = new SessionHolder(defaults);
     ServerCapabilities capabilities = new ServerCapabilities(
         new ServerCapabilities.Tools(null), null, null, null, null, null, null, Map.of());
-    Server server = new Server(new Implementation("rea-jadx-bridge", "1", null, null, null),
+    Server server = new Server(new Implementation("rea-jadx-bridge", "4", null, null, null),
         new ServerOptions(capabilities), "", ignored -> Unit.INSTANCE);
     SessionToolsKt.registerSessionTools(server, holder);
     ManifestToolsKt.registerManifestTools(server, holder);
@@ -187,13 +436,17 @@ class ReaJadxBridge {
     MethodToolsKt.registerMethodTools(server, holder);
     XrefToolsKt.registerXrefTools(server, holder);
     ReaJadxBridge bridge = new ReaJadxBridge(holder);
+    Tool loadTool = server.getTools().get("load_apk").getTool();
+    server.removeTool("load_apk");
+    server.addTool(loadTool, (connection, request, continuation) ->
+        bridge.load(request.getArguments(), defaults, continuation));
     Set<String> consumed = Set.of("load_apk", "get_app_info", "get_android_manifest",
         "list_classes", "get_class_summary", "get_method_body",
         "get_xrefs_to_class", "get_xrefs_to_method");
     for (String name : new ArrayList<>(server.getTools().keySet()))
       if (!consumed.contains(name)) server.removeTool(name);
     Tool diagnostics = new Tool("rea_jvm_status",
-        new ToolSchema(null, new JsonObject(Map.of()), List.of(), null),
+        new ToolSchema(null, new kotlinx.serialization.json.JsonObject(Map.of()), List.of(), null),
         "Observed worker JVM and metadata scope", null, null, null, null, null, null);
     server.addTool(diagnostics, (connection, request, continuation) -> bridge.handle("rea_jvm_status", request.getArguments()));
     for (String name : List.of("list_classes", "get_class_summary",
@@ -203,18 +456,21 @@ class ReaJadxBridge {
       List<String> fields = switch (name) {
         case "list_classes" -> List.of("offset", "limit");
         case "get_method_body" -> List.of("class_name", "method_name", "overload_index");
-        case "get_xrefs_to_method" -> List.of("class_name", "method_name");
+        case "get_xrefs_to_method" -> List.of("class_name", "method_name", "overload_index");
         default -> List.of("class_name");
       };
       Map<String, kotlinx.serialization.json.JsonElement> properties = new LinkedHashMap<>();
       for (String field : fields) properties.put(field, tool.getInputSchema().getProperties().get(field));
+      if (name.equals("list_classes")) properties.put("query",
+          new kotlinx.serialization.json.JsonObject(Map.of(
+              "type", kotlinx.serialization.json.JsonElementKt.JsonPrimitive("string"))));
       String description = switch (name) {
-        case "list_classes" -> "Indexed pages of all top-level class names, with offset and limit.";
+        case "list_classes" -> "Pages of case-sensitive matching top-level class names; offset and next_offset are inventory positions, limit bounds matching rows. Omitted or empty query inventories all names.";
         case "get_class_summary" -> "Parsed member metadata for one exact class name; no source generation.";
-        case "get_method_body" -> "Source for an exact class and stable same-name overload index, with smali fallback.";
-        default -> "Complete static references to an exact class or unique method; source lines are unresolved.";
+        case "get_method_body" -> "Source for an exact class and stable same-name overload index; static initializers use smali to preserve folded field initialization.";
+        default -> "Complete static using-node inventory for an exact class or selected method overload; one node may contain multiple invoke instructions, and source lines are unresolved.";
       };
-      Tool narrowed = new Tool(name, new ToolSchema(null, new JsonObject(properties),
+      Tool narrowed = new Tool(name, new ToolSchema(null, new kotlinx.serialization.json.JsonObject(properties),
           tool.getInputSchema().getRequired(), null), description, null, null, null, null, null, null);
       server.addTool(narrowed, (connection, request, continuation) -> {
         try { return bridge.handle(name, request.getArguments()); }