Convert files to data URLs on webhook post

python/cog/server/runner.py

@@ -16,7 +16,7 @@
 
 from .. import schema
 from ..base_input import BaseInput
-from ..files import put_file_to_signed_endpoint
+from ..files import put_file_to_signed_endpoint, upload_file
 from ..json import upload_files
 from ..types import PYDANTIC_V2
 from .errors import FileUploadError, RunnerBusyError, UnknownPredictionError
@@ -450,7 +450,7 @@ def _send_webhook(self, event: schema.WebhookEvent) -> None:
 
     def _upload_files(self, output: Any) -> Any:
         if self._file_uploader is None:
-            return output
+            return upload_file(output)
 
         try:
             # TODO: clean up output files

python/cog/server/webhook.py

@@ -47,11 +47,12 @@ def webhook_caller(webhook: str) -> Callable[[Any], None]:
     def caller(response: PredictionResponse) -> None:
         if throttler.should_send_response(response):
             if PYDANTIC_V2:
-                dict_response = jsonable_encoder(
-                    response.model_dump(exclude_unset=True)
-                )
+                response_object = response.model_dump(exclude_unset=True)
             else:
-                dict_response = jsonable_encoder(response.dict(exclude_unset=True))
+                response_object = response.dict(exclude_unset=True)
+
+            dict_response = jsonable_encoder(response_object)
+
             if Status.is_terminal(response.status):
                 # For terminal updates, retry persistently
                 retry_session.post(webhook, json=dict_response)

python/tests/server/test_webhook.py

@@ -1,9 +1,12 @@
+import tempfile
+
 import requests
 import responses
 from responses import registries
 
 from cog.schema import PredictionResponse, Status, WebhookEvent
 from cog.server.webhook import webhook_caller, webhook_caller_filtered
+from cog.types import Path
 
 
 @responses.activate
@@ -153,3 +156,28 @@ def test_webhook_caller_connection_errors():
     c = webhook_caller("https://example.com/webhook/123")
     # this should not raise an error
     c(response)
+
+
+@responses.activate
+def test_webhook_caller_converts_files_data_urls():
+    with tempfile.NamedTemporaryFile() as tmpfile:
+        with open(tmpfile.name, "w", encoding="utf8") as handle:
+            handle.write("hello world")
+
+        c = webhook_caller("https://example.com/webhook/123")
+
+        payload = {
+            "status": Status.SUCCEEDED,
+            "output": Path(tmpfile.name),
+            "input": {},
+        }
+        response = PredictionResponse(**payload)
+        payload["output"] = "data:None;base64,aGVsbG8gd29ybGQ="
+
+        responses.post(
+            "https://example.com/webhook/123",
+            json=payload,
+            status=200,
+        )
+
+        c(response)