Poll activities in their own organization

.changeset/poll-activity-org.yml

@@ -0,0 +1,4 @@
+type: patch
+packages: [ http ]
+changelog: |-
+  Activity polling now queries the organization of the submitted activity. Before, it used the client's default organization, so polls for sub-organization activities failed.

codegen/http/generate_http.py

@@ -231,6 +231,19 @@ def _request(self, url: str, body: Dict[str, Any], response_type: type) -> Any:
         response_data = response.json()
         return response_type(**response_data)
 
+    def _poll_body(self, activity: Any) -> GetActivityBody:
+        \"\"\"Build the get_activity body to poll an activity in its own organization.
+
+        Args:
+            activity: Activity object with id and organizationId attributes.
+
+        Returns:
+            GetActivityBody for the activity's organization. Falls back to the
+            client's default organization if the activity has none.
+        \"\"\"
+        organization_id = getattr(activity, "organizationId", None) or self.organization_id
+        return GetActivityBody(organizationId=organization_id, activityId=activity.id)
+
     def _poll_for_completion(self, activity: Any) -> Any:
         \"\"\"Poll until activity reaches terminal status.
 
@@ -246,7 +259,7 @@ def _poll_for_completion(self, activity: Any) -> Any:
         attempts = 0
         while attempts < self.max_polling_retries:
             time.sleep(self.polling_interval_ms / 1000.0)
-            poll_response = self.get_activity(GetActivityBody(activityId=activity.id))
+            poll_response = self.get_activity(self._poll_body(activity))
             activity = poll_response.activity
             if activity.status in TERMINAL_ACTIVITY_STATUSES:
                 break
@@ -352,7 +365,7 @@ def _poll_while_status(
             if time.monotonic() >= deadline:
                 return activity
 
-            poll_response = self.get_activity(GetActivityBody(activityId=activity.id))
+            poll_response = self.get_activity(self._poll_body(activity))
             activity = poll_response.activity
             status = getattr(activity.status, "value", activity.status)
             if status != watched_status:

packages/http/src/turnkey_http/generated/client.py

@@ -147,6 +147,21 @@ def _request(self, url: str, body: Dict[str, Any], response_type: type) -> Any:
         response_data = response.json()
         return response_type(**response_data)
 
+    def _poll_body(self, activity: Any) -> GetActivityBody:
+        """Build the get_activity body to poll an activity in its own organization.
+
+        Args:
+            activity: Activity object with id and organizationId attributes.
+
+        Returns:
+            GetActivityBody for the activity's organization. Falls back to the
+            client's default organization if the activity has none.
+        """
+        organization_id = (
+            getattr(activity, "organizationId", None) or self.organization_id
+        )
+        return GetActivityBody(organizationId=organization_id, activityId=activity.id)
+
     def _poll_for_completion(self, activity: Any) -> Any:
         """Poll until activity reaches terminal status.
 
@@ -162,7 +177,7 @@ def _poll_for_completion(self, activity: Any) -> Any:
         attempts = 0
         while attempts < self.max_polling_retries:
             time.sleep(self.polling_interval_ms / 1000.0)
-            poll_response = self.get_activity(GetActivityBody(activityId=activity.id))
+            poll_response = self.get_activity(self._poll_body(activity))
             activity = poll_response.activity
             if activity.status in TERMINAL_ACTIVITY_STATUSES:
                 break
@@ -270,7 +285,7 @@ def _poll_while_status(
             if time.monotonic() >= deadline:
                 return activity
 
-            poll_response = self.get_activity(GetActivityBody(activityId=activity.id))
+            poll_response = self.get_activity(self._poll_body(activity))
             activity = poll_response.activity
             status = getattr(activity.status, "value", activity.status)
             if status != watched_status:

packages/http/tests/test_activity_errors.py

@@ -5,6 +5,7 @@
 ``conftest.py``.
 """
 
+import json
 from typing import Any, Dict, List, Optional
 
 import pytest
@@ -25,6 +26,7 @@
 
 ACTIVITY_ID = "activity-123"
 ORG_ID = "org-123"
+SUB_ORG_ID = "sub-org-456"
 
 
 class StubStamper:
@@ -49,10 +51,11 @@ def activity_payload(
     result: Optional[Dict[str, Any]] = None,
     failure: Optional[Dict[str, Any]] = None,
     activity_type: str = "ACTIVITY_TYPE_CREATE_WALLET",
+    organization_id: str = ORG_ID,
 ) -> Dict[str, Any]:
     activity: Dict[str, Any] = {
         "id": ACTIVITY_ID,
-        "organizationId": ORG_ID,
+        "organizationId": organization_id,
         "status": status,
         "type": activity_type,
         "intent": {},
@@ -553,3 +556,62 @@ def test_raw_path_unpacks_versioned_result_key(client, mock_post):
 
     assert payload["subOrganizationId"] == "sub-1"
     assert payload["activity"]["status"] == "ACTIVITY_STATUS_COMPLETED"
+
+
+def poll_bodies(calls: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
+    return [
+        json.loads(call["data"])
+        for call in calls
+        if call["url"].endswith("/public/v1/query/get_activity")
+    ]
+
+
+def create_sub_org_wallet(client: TurnkeyClient):
+    return client.create_wallet(
+        CreateWalletBody(organizationId=SUB_ORG_ID, walletName="w", accounts=[])
+    )
+
+
+def completed_sub_org_wallet_payload() -> Dict[str, Any]:
+    return activity_payload(
+        "ACTIVITY_STATUS_COMPLETED",
+        result={"createWalletResult": {"walletId": "w-1", "addresses": ["0xabc"]}},
+        organization_id=SUB_ORG_ID,
+    )
+
+
+def test_pending_poll_uses_activity_organization(client, mock_post):
+    calls = mock_post(
+        activity_payload("ACTIVITY_STATUS_PENDING", organization_id=SUB_ORG_ID),
+        completed_sub_org_wallet_payload(),
+    )
+
+    response = create_sub_org_wallet(client)
+
+    assert response.walletId == "w-1"
+    assert json.loads(calls[0]["data"])["organizationId"] == SUB_ORG_ID
+    bodies = poll_bodies(calls)
+    assert len(bodies) == 1
+    assert bodies[0] == {"organizationId": SUB_ORG_ID, "activityId": ACTIVITY_ID}
+
+
+@pytest.mark.parametrize(
+    ("status", "enable"),
+    [
+        ("ACTIVITY_STATUS_AUTHENTICATORS_NEEDED", enable_mfa_polling),
+        ("ACTIVITY_STATUS_CONSENSUS_NEEDED", enable_consensus_polling),
+    ],
+)
+def test_approval_poll_uses_activity_organization(client, mock_post, status, enable):
+    calls = mock_post(
+        activity_payload(status, organization_id=SUB_ORG_ID),
+        completed_sub_org_wallet_payload(),
+    )
+    enable(client)
+
+    response = create_sub_org_wallet(client)
+
+    assert response.walletId == "w-1"
+    bodies = poll_bodies(calls)
+    assert len(bodies) == 1
+    assert bodies[0]["organizationId"] == SUB_ORG_ID