diff --git a/extensions/src/main/java/dev/cel/extensions/BUILD.bazel b/extensions/src/main/java/dev/cel/extensions/BUILD.bazel
index 8b7991cc0..7278b42e7 100644
--- a/extensions/src/main/java/dev/cel/extensions/BUILD.bazel
+++ b/extensions/src/main/java/dev/cel/extensions/BUILD.bazel
@@ -274,6 +274,7 @@ java_library(
         "//common/ast",
         "//common/internal:comparison_functions",
         "//common/types",
+        "//common/types:type_providers",
         "//compiler:compiler_builder",
         "//extensions:extension_library",
         "//parser:macro",
diff --git a/extensions/src/main/java/dev/cel/extensions/CelListsExtensions.java b/extensions/src/main/java/dev/cel/extensions/CelListsExtensions.java
index 79539b008..2dab1d558 100644
--- a/extensions/src/main/java/dev/cel/extensions/CelListsExtensions.java
+++ b/extensions/src/main/java/dev/cel/extensions/CelListsExtensions.java
@@ -18,6 +18,7 @@
 import static com.google.common.base.Preconditions.checkNotNull;
 import static com.google.common.collect.ImmutableSet.toImmutableSet;
 
+import com.google.common.base.Ascii;
 import com.google.common.base.Preconditions;
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableSet;
@@ -31,6 +32,7 @@
 import dev.cel.common.Operator;
 import dev.cel.common.ast.CelExpr;
 import dev.cel.common.internal.ComparisonFunctions;
+import dev.cel.common.types.CelType;
 import dev.cel.common.types.ListType;
 import dev.cel.common.types.SimpleType;
 import dev.cel.common.types.TypeParamType;
@@ -54,6 +56,10 @@
 public final class CelListsExtensions
     implements CelCompilerLibrary, CelInternalRuntimeLibrary, CelExtensionLibrary.FeatureSet {
 
+  private static final CelObjectComparator OBJECT_COMPARATOR = new CelObjectComparator();
+  private static final String UNUSED_ITER_VAR = "#unused";
+  private static final String SORT_BY_INPUT_VAR = "@__sortBy_input__";
+
   /** Supported functions for Lists extension library. */
   @SuppressWarnings({"unchecked"}) // Unchecked: Type-checker guarantees casting safety.
   public enum Function {
@@ -131,17 +137,50 @@ public enum Function {
                 ListType.create(TypeParamType.create("T")))),
         CelFunctionBinding.from("list_sort", Collection.class, CelListsExtensions::sort)),
     SORT_BY(
-        CelFunctionDecl.newFunctionDeclaration(
-            "lists.@sortByAssociatedKeys",
-            CelOverloadDecl.newGlobalOverload(
-                "list_sortByAssociatedKeys",
-                "Sorts a list by a key value. Used by the 'sortBy' macro",
+        createSortByFunctionDecl(comparableSortKeyTypes()),
+        createSortByFunctionBindings(comparableSortKeyTypes()));
+
+    private static ImmutableList<CelType> comparableSortKeyTypes() {
+      return ImmutableList.of(
+          SimpleType.INT,
+          SimpleType.UINT,
+          SimpleType.DOUBLE,
+          SimpleType.BOOL,
+          SimpleType.STRING,
+          SimpleType.BYTES,
+          SimpleType.DURATION,
+          SimpleType.TIMESTAMP);
+    }
+
+    private static CelFunctionDecl createSortByFunctionDecl(ImmutableList<CelType> keyTypes) {
+      ImmutableList.Builder<CelOverloadDecl> overloads = ImmutableList.builder();
+      for (CelType type : keyTypes) {
+        String typeName = Ascii.toLowerCase(type.kind().name());
+        overloads.add(
+            CelOverloadDecl.newMemberOverload(
+                String.format("list_%s_sortByAssociatedKeys", typeName),
+                "Sorts a list by an associated list of keys. Used by the 'sortBy' macro",
                 ListType.create(TypeParamType.create("T")),
-                ListType.create(TypeParamType.create("T")))),
-        CelFunctionBinding.from(
-            "list_sortByAssociatedKeys",
-            Collection.class,
-            CelListsExtensions::sortByAssociatedKeys));
+                ListType.create(TypeParamType.create("T")),
+                ListType.create(type)));
+      }
+      return CelFunctionDecl.newFunctionDeclaration("@sortByAssociatedKeys", overloads.build());
+    }
+
+    private static CelFunctionBinding[] createSortByFunctionBindings(
+        ImmutableList<CelType> keyTypes) {
+      return keyTypes.stream()
+          .map(
+              type -> {
+                String typeName = Ascii.toLowerCase(type.kind().name());
+                return CelFunctionBinding.from(
+                    String.format("list_%s_sortByAssociatedKeys", typeName),
+                    Collection.class,
+                    Collection.class,
+                    CelListsExtensions::sortByAssociatedKeys);
+              })
+          .toArray(CelFunctionBinding[]::new);
+    }
 
     private final CelFunctionDecl functionDecl;
     private final ImmutableSet<CelFunctionBinding> functionBindings;
@@ -359,7 +398,15 @@ private static List<Object> reverse(Collection<Object> list) {
   }
 
   private static ImmutableList<Object> sort(Collection<Object> objects) {
-    return ImmutableList.sortedCopyOf(new CelObjectComparator(), objects);
+    if (objects.isEmpty()) {
+      return ImmutableList.of();
+    }
+    if (objects.size() == 1) {
+      Object single = objects.iterator().next();
+      OBJECT_COMPARATOR.compare(single, single);
+      return ImmutableList.of(single);
+    }
+    return ImmutableList.sortedCopyOf(OBJECT_COMPARATOR, objects);
   }
 
   private static class CelObjectComparator implements Comparator<Object> {
@@ -383,6 +430,34 @@ public int compare(Object o1, Object o2) {
     }
   }
 
+  /**
+   * Expands the {@code list.sortBy(var, expr)} receiver macro into a binding expression that sorts
+   * the target list using keys evaluated by mapping {@code expr} over each element.
+   *
+   * <p>For example, given:
+   *
+   * <pre>{@code
+   * myList.sortBy(item, -item.field)
+   * }</pre>
+   *
+   * <p>The macro expands into:
+   *
+   * <pre>{@code
+   * cel.bind(@__sortBy_input__, myList,
+   *     @__sortBy_input__.@sortByAssociatedKeys(
+   *         @__sortBy_input__.map(item, -item.field)
+   *     )
+   * )
+   * }</pre>
+   *
+   * <p>Where:
+   *
+   * <ul>
+   *   <li>{@code @__sortBy_input__.map(item, -item.field)} evaluates the sort key for each element.
+   *   <li>{@code @sortByAssociatedKeys} stably sorts the input list elements based on their
+   *       corresponding sort keys.
+   * </ul>
+   */
   private static Optional<CelExpr> sortByMacro(
       CelMacroExprFactory exprFactory, CelExpr target, ImmutableList<CelExpr> arguments) {
     checkNotNull(exprFactory);
@@ -400,56 +475,86 @@ private static Optional<CelExpr> sortByMacro(
     String varName = varIdent.ident().name();
     CelExpr sortKeyExpr = checkNotNull(arguments.get(1));
 
-    // Compute the key using the second argument of the `sortBy(e, key)` macro.
-    // Combine the key and the value in a two-element list
-    CelExpr step = exprFactory.newList(sortKeyExpr, varIdent);
-    // Wrap the pair in another list in order to be able to use the `list+list` operator
-    step = exprFactory.newList(step);
-    // Append the key-value pair to the i
-    step =
+    // Build map comprehension: @__sortBy_input__.map(varName, sortKeyExpr)
+    CelExpr targetIdent = exprFactory.newIdentifier(SORT_BY_INPUT_VAR);
+    CelExpr mapStep =
         exprFactory.newGlobalCall(
             Operator.ADD.getFunction(),
             exprFactory.newIdentifier(exprFactory.getAccumulatorVarName()),
-            step);
-    // Create an intermediate list and populate it with key-value pairs
-    step =
+            exprFactory.newList(sortKeyExpr));
+    CelExpr mapCompr =
         exprFactory.fold(
             varName,
-            target,
+            targetIdent,
             exprFactory.getAccumulatorVarName(),
             exprFactory.newList(),
-            exprFactory.newBoolLiteral(true), // Include all elements
-            step,
+            exprFactory.newBoolLiteral(true),
+            mapStep,
             exprFactory.newIdentifier(exprFactory.getAccumulatorVarName()));
-    // Finally, sort the list of key-value pairs and map it to a list of values
-    step = exprFactory.newGlobalCall(Function.SORT_BY.getFunction(), step);
 
-    return Optional.of(step);
+    // Build call: @__sortBy_input__.@sortByAssociatedKeys(mapCompr)
+    CelExpr callExpr =
+        exprFactory.newReceiverCall(
+            Function.SORT_BY.getFunction(), exprFactory.newIdentifier(SORT_BY_INPUT_VAR), mapCompr);
+
+    // Build bind: cel.bind(@__sortBy_input__, target, callExpr)
+    CelExpr bindExpr =
+        exprFactory.fold(
+            UNUSED_ITER_VAR,
+            exprFactory.newList(),
+            SORT_BY_INPUT_VAR,
+            target,
+            exprFactory.newBoolLiteral(false),
+            exprFactory.newIdentifier(SORT_BY_INPUT_VAR),
+            callExpr);
+
+    return Optional.of(bindExpr);
   }
 
-  @SuppressWarnings({"unchecked", "rawtypes"})
+  /**
+   * Sorts elements of {@code list} based on the natural order of corresponding elements in {@code
+   * keys}.
+   *
+   * <p>Both {@code list} and {@code keys} must have the exact same size. The sorting is stable
+   * (i.e., preserves the relative order of elements with equal keys).
+   *
+   * @param list The input list to sort
+   * @param keys The associated keys evaluated for each element in {@code list}
+   * @return A new {@link ImmutableList} containing the elements of {@code list} sorted by {@code
+   *     keys}
+   */
   private static ImmutableList<Object> sortByAssociatedKeys(
-      Collection<List<Object>> keyValuePairs) {
-    List<Object>[] array = keyValuePairs.toArray(new List[0]);
-    Arrays.sort(array, new CelObjectByKeyComparator(new CelObjectComparator()));
-    ImmutableList.Builder<Object> builder = ImmutableList.builderWithExpectedSize(array.length);
-    for (List<Object> pair : array) {
-      builder.add(pair.get(1));
+      Collection<Object> list, Collection<Object> keys) {
+    checkArgument(
+        list.size() == keys.size(),
+        "@sortByAssociatedKeys() expected a list of the same size as the associated keys"
+            + " list, but got %s in list and %s in keys",
+        list.size(),
+        keys.size());
+
+    int listSize = list.size();
+    if (listSize == 0) {
+      return ImmutableList.of();
     }
-    return builder.build();
-  }
 
-  private static class CelObjectByKeyComparator implements Comparator<Object> {
-    private final CelObjectComparator keyComparator;
+    Object[] listArray = list.toArray();
+    Object[] keysArray = keys.toArray();
+    if (listSize == 1) {
+      OBJECT_COMPARATOR.compare(keysArray[0], keysArray[0]);
+      return ImmutableList.of(listArray[0]);
+    }
 
-    CelObjectByKeyComparator(CelObjectComparator keyComparator) {
-      this.keyComparator = keyComparator;
+    Integer[] indices = new Integer[listSize];
+    for (int i = 0; i < listSize; i++) {
+      indices[i] = i;
     }
 
-    @SuppressWarnings({"unchecked"})
-    @Override
-    public int compare(Object o1, Object o2) {
-      return keyComparator.compare(((List<Object>) o1).get(0), ((List<Object>) o2).get(0));
+    Arrays.sort(indices, (i1, i2) -> OBJECT_COMPARATOR.compare(keysArray[i1], keysArray[i2]));
+
+    ImmutableList.Builder<Object> builder = ImmutableList.builderWithExpectedSize(listSize);
+    for (int index : indices) {
+      builder.add(listArray[index]);
     }
+    return builder.build();
   }
 }
diff --git a/extensions/src/test/java/dev/cel/extensions/BUILD.bazel b/extensions/src/test/java/dev/cel/extensions/BUILD.bazel
index 920ba537b..f7b996610 100644
--- a/extensions/src/test/java/dev/cel/extensions/BUILD.bazel
+++ b/extensions/src/test/java/dev/cel/extensions/BUILD.bazel
@@ -42,12 +42,14 @@ java_library(
         "//parser:unparser",
         "//runtime",
         "//runtime:function_binding",
-        "//runtime:interpreter_util",
         "//runtime:lite_runtime",
         "//runtime:lite_runtime_factory",
         "//runtime:partial_vars",
         "//runtime:unknown_attributes",
         "//testing:cel_runtime_flavor",
+        "//validator",
+        "//validator:validator_builder",
+        "//validator/validators:homogeneous_literal",
         "@cel_spec//proto/cel/expr/conformance/proto2:test_all_types_java_proto",
         "@cel_spec//proto/cel/expr/conformance/proto3:test_all_types_java_proto",
         "@cel_spec//proto/cel/expr/conformance/test:simple_java_proto",
diff --git a/extensions/src/test/java/dev/cel/extensions/CelExtensionsTest.java b/extensions/src/test/java/dev/cel/extensions/CelExtensionsTest.java
index 31c7d65c8..279ad7013 100644
--- a/extensions/src/test/java/dev/cel/extensions/CelExtensionsTest.java
+++ b/extensions/src/test/java/dev/cel/extensions/CelExtensionsTest.java
@@ -185,7 +185,7 @@ public void getAllFunctionNames() {
             "distinct",
             "reverse",
             "sort",
-            "lists.@sortByAssociatedKeys",
+            "@sortByAssociatedKeys",
             "regex.replace",
             "regex.extract",
             "regex.extractAll",
diff --git a/extensions/src/test/java/dev/cel/extensions/CelListsExtensionsTest.java b/extensions/src/test/java/dev/cel/extensions/CelListsExtensionsTest.java
index 4520f81ba..64d281422 100644
--- a/extensions/src/test/java/dev/cel/extensions/CelListsExtensionsTest.java
+++ b/extensions/src/test/java/dev/cel/extensions/CelListsExtensionsTest.java
@@ -22,6 +22,7 @@
 import com.google.testing.junit.testparameterinjector.TestParameterInjector;
 import com.google.testing.junit.testparameterinjector.TestParameters;
 import dev.cel.bundle.Cel;
+import dev.cel.common.CelAbstractSyntaxTree;
 import dev.cel.common.CelContainer;
 import dev.cel.common.CelValidationException;
 import dev.cel.common.CelValidationResult;
@@ -30,6 +31,9 @@
 import dev.cel.parser.CelStandardMacro;
 import dev.cel.runtime.CelEvaluationException;
 import dev.cel.testing.CelRuntimeFlavor;
+import dev.cel.validator.CelValidator;
+import dev.cel.validator.CelValidatorFactory;
+import dev.cel.validator.validators.HomogeneousLiteralValidator;
 import org.junit.Assume;
 import org.junit.Test;
 import org.junit.runner.RunWith;
@@ -64,7 +68,7 @@ public void functionList_byVersion() {
             "distinct",
             "reverse",
             "sort",
-            "lists.@sortByAssociatedKeys");
+            "@sortByAssociatedKeys");
   }
 
   @Test
@@ -257,6 +261,9 @@ public void sort_success_heterogeneousNumbers(String expression, String expected
   @TestParameters(
       "{expression: '[SimpleTest{name: \"a\"}, SimpleTest{name: \"b\"}].sort()', "
           + "expectedError: 'List elements must be comparable'}")
+  @TestParameters(
+      "{expression: '[SimpleTest{name: \"a\"}].sort()', "
+          + "expectedError: 'List elements must be comparable'}")
   public void sort_throws(String expression, String expectedError) throws Exception {
     assertThat(assertThrows(CelEvaluationException.class, () -> eval(cel, expression)))
         .hasCauseThat()
@@ -283,6 +290,11 @@ public void sort_throws(String expression, String expectedError) throws Exceptio
           + "expected: '[SimpleTest{name: \"bar\"},"
           + " SimpleTest{name: \"baz\"},"
           + " SimpleTest{name: \"foo\"}]'}")
+  @TestParameters(
+      "{expression: '[SimpleTest{name: \"baz\"},"
+          + " SimpleTest{name: \"foo\"},"
+          + " SimpleTest{name: \"bar\"}].sortBy(e, e.name)[0].name', "
+          + "expected: '\"bar\"'}")
   public void sortBy_success(String expression, String expected) throws Exception {
     Object result = eval(cel, expression);
 
@@ -296,6 +308,12 @@ public void sortBy_success(String expression, String expected) throws Exception
   @TestParameters(
       "{expression: 'lists.range(3).sortBy(e.foo, e)', "
           + "expectedError: 'variable name must be a simple identifier'}")
+  @TestParameters(
+      "{expression: '[SimpleTest{name: \"a\"}].sortBy(e, e)', "
+          + "expectedError: 'found no matching overload for ''@sortByAssociatedKeys'''}")
+  @TestParameters(
+      "{expression: '[SimpleTest{name: \"a\"}, SimpleTest{name: \"b\"}].sortBy(e, e)', "
+          + "expectedError: 'found no matching overload for ''@sortByAssociatedKeys'''}")
   public void sortBy_throws_validationException(String expression, String expectedError)
       throws Exception {
     CelValidationResult result = cel.compile(expression);
@@ -305,19 +323,20 @@ public void sortBy_throws_validationException(String expression, String expected
   }
 
   @Test
-  @TestParameters(
-      "{expression: '[[1, 2], [\"a\", \"b\"]].sortBy(e, e[0])', "
-          + "expectedError: 'List elements must have the same type'}")
-  @TestParameters(
-      "{expression: '[SimpleTest{name: \"a\"}, SimpleTest{name: \"b\"}].sortBy(e, e)', "
-          + "expectedError: 'List elements must be comparable'}")
-  public void sortBy_throws_evaluationException(String expression, String expectedError)
-      throws Exception {
-    assertThat(assertThrows(CelEvaluationException.class, () -> eval(cel, expression)))
-        .hasCauseThat()
-        .hasMessageThat()
-        .contains(expectedError);
+  public void sortBy_withHomogeneousLiteralValidator_success() throws Exception {
+    CelValidator validator =
+        CelValidatorFactory.standardCelValidatorBuilder(cel)
+            .addAstValidators(HomogeneousLiteralValidator.newInstance())
+            .build();
+
+    CelAbstractSyntaxTree ast =
+        cel.compile(
+                "[SimpleTest{name: 'baz'}, SimpleTest{name: 'foo'}, SimpleTest{name: 'bar'}]"
+                    + ".sortBy(e, e.name)[0].name")
+            .getAst();
+    CelValidationResult result = validator.validate(ast);
+
+    assertThat(result.hasError()).isFalse();
+    assertThat(cel.createProgram(ast).eval()).isEqualTo("bar");
   }
-
-
 }