From f94a4e6ab606c5d3aeb14a6becc46b5756581ee0 Mon Sep 17 00:00:00 2001
From: Laura Abro <laura@koii.network>
Date: Thu, 24 Apr 2025 12:09:02 -0300
Subject: [PATCH] add worker submission

---
 worker/orca-agent/tests/stages/worker_pr.py   | 11 +++-
 .../tests/stages/worker_submission.py         | 60 +++++++++++++++++++
 2 files changed, 68 insertions(+), 3 deletions(-)

diff --git a/worker/orca-agent/tests/stages/worker_pr.py b/worker/orca-agent/tests/stages/worker_pr.py
index 51cc7bf..ff61866 100644
--- a/worker/orca-agent/tests/stages/worker_pr.py
+++ b/worker/orca-agent/tests/stages/worker_pr.py
@@ -4,6 +4,11 @@ from prometheus_test.utils import create_signature
 
 def prepare(runner, worker):
     round_state = runner.state["rounds"][str(runner.current_round)]
+
+    if worker.name not in round_state["pr_urls"]:
+        print(f"✓ No PR URL found for {worker.name} - continuing")
+        return None
+
     payload = {
         "taskId": runner.config.task_id,
         "action": "add-todo-pr",
@@ -20,9 +25,9 @@ def prepare(runner, worker):
 
 def execute(runner, worker, data):
     """Add worker PR URL to middle server"""
-    if not runner.state["pr_url"]:
-        print(f"✓ No PR URL found for {worker.name} - continuing")
-        return {"success": True, "message": "No PR URL found"}
+
+    if data is None:
+        return {"success": True, "message": "Skipped due to missing PR URL"}
 
     url = f"{runner.config.middle_server_url}/summarizer/worker/add-todo-pr"
     response = requests.post(
diff --git a/worker/orca-agent/tests/stages/worker_submission.py b/worker/orca-agent/tests/stages/worker_submission.py
index e69de29..39b0930 100644
--- a/worker/orca-agent/tests/stages/worker_submission.py
+++ b/worker/orca-agent/tests/stages/worker_submission.py
@@ -0,0 +1,60 @@
+"""Stage for handling worker submissions."""
+
+import requests
+from prometheus_test.utils import create_signature
+
+
+def prepare(runner, worker):
+    """Prepare data for worker submission"""
+    # Get the current round's state
+    round_state = runner.state.get("rounds", {}).get(str(runner.current_round), {})
+    pr_urls = round_state.get("pr_urls", {})
+
+    if worker.name not in pr_urls:
+        # Return None to indicate this step should be skipped
+        print(f"✓ No PR URL found for {worker.name} - continuing")
+        return None
+
+    # Get submission data from worker
+    url = f"{worker.url}/submission/{runner.current_round}"
+    response = requests.get(url)
+    response.raise_for_status()
+    submission_data = response.json()
+
+    # Create signature for the submission
+    submitter_payload = {
+        "taskId": runner.config.task_id,
+        "roundNumber": runner.current_round,
+        "stakingKey": worker.staking_public_key,
+        "pubKey": worker.public_key,
+        "action": "audit",
+        **submission_data,
+    }
+
+    return {
+        **submission_data,
+        "signature": create_signature(worker.staking_signing_key, submitter_payload),
+        "stakingKey": worker.staking_public_key,
+        "pubKey": worker.public_key,
+    }
+
+
+def execute(runner, worker, data):
+    """Store worker submission data"""
+    # If prepare returned None, skip this step
+    if data is None:
+        return {"success": True, "message": "Skipped due to missing PR URL"}
+
+    # Store submission data in state
+    round_key = str(runner.current_round)
+    round_state = runner.state["rounds"].setdefault(round_key, {})
+
+    # Initialize submission_data if not exists
+    if "submission_data" not in round_state:
+        round_state["submission_data"] = {}
+
+    # Store or update submission data
+    round_state["submission_data"][worker.name] = data
+
+    # Return success result
+    return {"success": True, "data": data}