From 5f9511527e6284bb508c7a4c846e3b5f4f733748 Mon Sep 17 00:00:00 2001
From: Abhinav Kumar Singh <abhinav.kr.singh.2610@gmail.com>
Date: Sat, 15 Aug 2026 14:46:51 +0530
Subject: [PATCH] fix: publish ARM64 local-mode wheels

---
 .github/workflows/publish-pypi.yml |  8 ++++++++
 tests/test_sea_binary.py           | 21 +++++++++++++++++++++
 2 files changed, 29 insertions(+)

diff --git a/.github/workflows/publish-pypi.yml b/.github/workflows/publish-pypi.yml
index 2f2e895b..67efe9e2 100644
--- a/.github/workflows/publish-pypi.yml
+++ b/.github/workflows/publish-pypi.yml
@@ -20,6 +20,10 @@ jobs:
             binary_name: stagehand-server-v3-linux-x64
             output_path: src/stagehand/_sea/stagehand-linux-x64
             wheel_platform_tag: manylinux2014_x86_64
+          - os: ubuntu-latest
+            binary_name: stagehand-server-v3-linux-arm64
+            output_path: src/stagehand/_sea/stagehand-linux-arm64
+            wheel_platform_tag: manylinux2014_aarch64
           - os: macos-latest
             binary_name: stagehand-server-v3-darwin-arm64
             output_path: src/stagehand/_sea/stagehand-darwin-arm64
@@ -32,6 +36,10 @@ jobs:
             binary_name: stagehand-server-v3-win32-x64.exe
             output_path: src/stagehand/_sea/stagehand-win32-x64.exe
             wheel_platform_tag: win_amd64
+          - os: windows-latest
+            binary_name: stagehand-server-v3-win32-arm64.exe
+            output_path: src/stagehand/_sea/stagehand-win32-arm64.exe
+            wheel_platform_tag: win_arm64
 
     runs-on: ${{ matrix.os }}
     permissions:
diff --git a/tests/test_sea_binary.py b/tests/test_sea_binary.py
index 1c098ac7..ea949872 100644
--- a/tests/test_sea_binary.py
+++ b/tests/test_sea_binary.py
@@ -23,6 +23,27 @@ def _load_download_binary_module():
 download_binary = _load_download_binary_module()
 
 
+@pytest.mark.parametrize(
+    ("sys_platform", "machine", "expected"),
+    [
+        ("linux", "aarch64", "stagehand-linux-arm64"),
+        ("linux", "x86_64", "stagehand-linux-x64"),
+        ("win32", "ARM64", "stagehand-win32-arm64.exe"),
+        ("win32", "AMD64", "stagehand-win32-x64.exe"),
+    ],
+)
+def test_default_binary_filename_matches_published_wheel(
+    monkeypatch: pytest.MonkeyPatch,
+    sys_platform: str,
+    machine: str,
+    expected: str,
+) -> None:
+    monkeypatch.setattr(sea_binary.sys, "platform", sys_platform)
+    monkeypatch.setattr(sea_binary.platform, "machine", lambda: machine)
+
+    assert sea_binary.default_binary_filename() == expected
+
+
 def test_resolve_binary_path_defaults_cache_version_to_package_version(
     monkeypatch: pytest.MonkeyPatch,
     tmp_path: Path,