From a3c3b9cc167153027bdfb4e6bdb70d8792d442c1 Mon Sep 17 00:00:00 2001
From: Michael Bunsen <notbot@gmail.com>
Date: Fri, 23 Jan 2026 23:06:11 -0800
Subject: [PATCH] Add error field to SourceImageResponse for unified error
 handling

- Add optional `error` field to SourceImageResponse
- Make `url` optional in SourceImageResponse (may be None for failed images)
- Remove AntennaTaskResultError - errors now go in source_images[0].error
- Simplify AntennaTaskResult to always use PipelineResultsResponse
- Update worker to create error responses with the new pattern

This provides a cleaner API where success/error is determined by checking
`result.source_images[0].error is None` rather than type checking.

Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>
---
 trapdata/api/schemas.py           | 22 ++++++++++++----------
 trapdata/api/tests/test_worker.py | 11 ++++-------
 trapdata/cli/worker.py            | 18 +++++++++++++-----
 3 files changed, 29 insertions(+), 22 deletions(-)

diff --git a/trapdata/api/schemas.py b/trapdata/api/schemas.py
index 30fd6186..c9515132 100644
--- a/trapdata/api/schemas.py
+++ b/trapdata/api/schemas.py
@@ -137,10 +137,16 @@ class SourceImageRequest(pydantic.BaseModel):
 
 
 class SourceImageResponse(pydantic.BaseModel):
+    """Response for a source image, with optional error if processing failed."""
+
     model_config = pydantic.ConfigDict(extra="ignore")
 
     id: str
-    url: str
+    url: str | None = None
+    error: str | None = pydantic.Field(
+        default=None,
+        description="Error message if the image failed to process. None means success.",
+    )
 
 
 class AlgorithmCategoryMapResponse(pydantic.BaseModel):
@@ -342,18 +348,14 @@ class PipelineConfigResponse(pydantic.BaseModel):
     stages: list[PipelineStage] = []
 
 
-class AntennaTaskResultError(pydantic.BaseModel):
-    """Error result for a single Antenna task that failed to process."""
-
-    error: str
-    image_id: str | None = None
-
-
 class AntennaTaskResult(pydantic.BaseModel):
-    """Result for a single Antenna task, either success or error."""
+    """Result for a single Antenna task.
+
+    Check source_images[0].error for failure. If error is None, the task succeeded.
+    """
 
     reply_subject: str | None = None
-    result: PipelineResultsResponse | AntennaTaskResultError
+    result: PipelineResultsResponse
 
 
 class AntennaTaskResults(pydantic.BaseModel):
diff --git a/trapdata/api/tests/test_worker.py b/trapdata/api/tests/test_worker.py
index f9cc07b2..ef55706c 100644
--- a/trapdata/api/tests/test_worker.py
+++ b/trapdata/api/tests/test_worker.py
@@ -10,11 +10,7 @@
 import torch
 
 from trapdata.api.datasets import RESTDataset, rest_collate_fn
-from trapdata.api.schemas import (
-    AntennaTaskResult,
-    AntennaTaskResultError,
-    PipelineResultsResponse,
-)
+from trapdata.api.schemas import AntennaTaskResult, PipelineResultsResponse
 from trapdata.cli.worker import _get_jobs, _process_job
 
 # ---------------------------------------------------------------------------
@@ -493,11 +489,12 @@ def test_handles_failed_items(self, mock_loader_fn, mock_detector_cls, mock_post
         batch_results = mock_post.call_args[0][2]
         # 1 success + 1 failure
         assert len(batch_results) == 2
+        # Find error items by checking source_images[0].error
         error_items = [
-            r for r in batch_results if isinstance(r.result, AntennaTaskResultError)
+            r for r in batch_results if r.result.source_images[0].error is not None
         ]
         assert len(error_items) == 1
-        assert error_items[0].result.error == "404 not found"
+        assert error_items[0].result.source_images[0].error == "404 not found"
         assert error_items[0].reply_subject == "reply_fail"
 
     @patch("trapdata.cli.worker.get_rest_dataloader")
diff --git a/trapdata/cli/worker.py b/trapdata/cli/worker.py
index 3ffb9a15..6956bbe6 100644
--- a/trapdata/cli/worker.py
+++ b/trapdata/cli/worker.py
@@ -14,7 +14,6 @@
 from trapdata.api.schemas import (
     AntennaJobsListResponse,
     AntennaTaskResult,
-    AntennaTaskResultError,
     DetectionResponse,
     PipelineResultsResponse,
     SourceImageResponse,
@@ -260,13 +259,22 @@ def _process_job(pipeline: str, job_id: int, settings: Settings) -> bool:
         failed_items = batch.get("failed_items")
         if failed_items:
             for failed_item in failed_items:
+                # Create error response with error in SourceImageResponse
+                error_source_image = SourceImageResponse(
+                    id=failed_item.get("image_id", "unknown"),
+                    url=failed_item.get("image_url"),
+                    error=failed_item.get("error", "Unknown error"),
+                )
+                error_response = PipelineResultsResponse(
+                    pipeline=pipeline,
+                    source_images=[error_source_image],
+                    detections=[],
+                    total_time=0.0,
+                )
                 batch_results.append(
                     AntennaTaskResult(
                         reply_subject=failed_item.get("reply_subject"),
-                        result=AntennaTaskResultError(
-                            error=failed_item.get("error", "Unknown error"),
-                            image_id=failed_item.get("image_id"),
-                        ),
+                        result=error_response,
                     )
                 )