diff --git a/bazel/graal/rules.bzl b/bazel/graal.bzl
similarity index 74%
rename from bazel/graal/rules.bzl
rename to bazel/graal.bzl
index 7d554b6a994..a052e2aab75 100644
--- a/bazel/graal/rules.bzl
+++ b/bazel/graal.bzl
@@ -17,29 +17,19 @@
 def _graal_native_binary_impl(ctx):
     cc_toolchain_info = ctx.toolchains["@bazel_tools//tools/cpp:toolchain_type"].cc
 
-    # TODO(james): this is a bit of a hack because bazel's JavaInfo doesn't currently have starlark definitions,
-    # so its hard to traverse.
-    default_info = ctx.attr.java_binary[DefaultInfo]
-    jar = None
-    for file in default_info.files.to_list():
-        if file.path.endswith(".jar"):
-            jar = file
-            break
-    if jar == None:
-        fail("no .jar file in java_binary rule output")
+    if len(ctx.attr.java_binary[JavaInfo].java_outputs) == 0:
+        fail("no jars in java_binary rule output")
 
-    out_name = ctx.attr.output_name
-    if out_name == "":
-        out_name = ctx.attr.name
-
-    out = ctx.actions.declare_file(out_name)
+    if len(ctx.attr.java_binary[JavaInfo].java_outputs) > 1:
+        fail("more than one output jar in java_binary rule output")
 
+    jar = ctx.attr.java_binary[JavaInfo].java_outputs[0].class_jar
     graal_runtime = ctx.attr.graal_runtime[java_common.JavaRuntimeInfo]
     args = [
         "-cp",
         jar.path,
         "-o",
-        out.path,
+        ctx.outputs.output_name.path,
         "--native-compiler-path=" + cc_toolchain_info.compiler_executable,
         # Add /usr/bin as prefix, so that `native-image` can find ld.
         # The real solution would be to get `native-image` to work with the combination of lld and gcc.
@@ -49,7 +39,7 @@ def _graal_native_binary_impl(ctx):
     ] + ctx.attr.extra_args
 
     ctx.actions.run(
-        outputs = [out],
+        outputs = [ctx.outputs.output_name],
         inputs = depset(
             [jar],
             transitive = [
@@ -63,8 +53,8 @@ def _graal_native_binary_impl(ctx):
 
     return [
         DefaultInfo(
-            files = depset([out]),
-            executable = out,
+            files = depset([ctx.outputs.output_name]),
+            executable = ctx.outputs.output_name,
         ),
     ]
 
@@ -79,9 +69,9 @@ graal_native_binary = rule(
             allow_files = True,
         ),
         "java_binary": attr.label(
-            providers = [JavaInfo, DefaultInfo],
+            providers = [JavaInfo],
         ),
-        "output_name": attr.string(),
+        "output_name": attr.output(mandatory = True),
     },
     toolchains = [
         "@bazel_tools//tools/cpp:toolchain_type",
diff --git a/bazel/graal/BUILD.bazel b/bazel/graal/BUILD.bazel
deleted file mode 100644
index 14ca481dc04..00000000000
--- a/bazel/graal/BUILD.bazel
+++ /dev/null
@@ -1,15 +0,0 @@
-# Copyright 2018- The Pixie Authors.
-#
-# 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
-#
-#     http://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.
-#
-# SPDX-License-Identifier: Apache-2.0
diff --git a/src/stirling/source_connectors/perf_profiler/testing/java/BUILD.bazel b/src/stirling/source_connectors/perf_profiler/testing/java/BUILD.bazel
index 02c08d6957f..3f963deed13 100644
--- a/src/stirling/source_connectors/perf_profiler/testing/java/BUILD.bazel
+++ b/src/stirling/source_connectors/perf_profiler/testing/java/BUILD.bazel
@@ -17,9 +17,9 @@
 load("@io_bazel_rules_docker//container:image.bzl", "container_image")
 load("@io_bazel_rules_docker//java:image.bzl", "java_image")
 load("@rules_java//java:defs.bzl", "java_binary")
+load("//bazel:graal.bzl", "graal_native_binary")
 load("//bazel:pl_build_system.bzl", "pl_cc_binary")
 load("//bazel:toolchain_transitions.bzl", "java_graal_binary")
-load("//bazel/graal:rules.bzl", "graal_native_binary")
 load("//src/stirling/source_connectors/perf_profiler/testing:testing.bzl", "jdk_names")
 
 package(default_visibility = ["//src/stirling:__subpackages__"])