Python: update ChatKit attachment handling

python/samples/05-end-to-end/chatkit-integration/app.py

@@ -43,7 +43,7 @@
 # ChatKit imports
 from chatkit.actions import Action
 from chatkit.server import ChatKitServer
-from chatkit.store import StoreItemType, default_generate_id
+from chatkit.store import NotFoundError, StoreItemType, default_generate_id
 from chatkit.types import (
     ThreadItem,
     ThreadItemDoneEvent,
@@ -595,20 +595,25 @@ async def upload_file(attachment_id: str, file: UploadFile = File(...)):  # noqa
         logger.warning(f"Rejected invalid attachment ID: {attachment_id!r}")
         return JSONResponse(status_code=400, content={"error": "Invalid attachment ID."})
 
+    try:
+        attachment = await data_store.load_attachment(attachment_id, {"user_id": DEFAULT_USER_ID})
+    except NotFoundError:
+        return JSONResponse(status_code=404, content={"error": "Attachment not found."})
+
+    if attachment.upload_descriptor is None:
+        return JSONResponse(status_code=409, content={"error": "Attachment upload is already complete."})
+
     try:
         # Read file contents
         contents = await file.read()
 
         # Save to disk
-        file_path.write_bytes(contents)
+        file_path.write_bytes(contents)  # CodeQL [SM01305] Path is constrained by get_file_path.
 
         logger.info(f"Saved {len(contents)} bytes to {file_path}")
 
-        # Load the attachment metadata from the data store
-        attachment = await data_store.load_attachment(attachment_id, {"user_id": DEFAULT_USER_ID})
-
-        # Clear the upload_url since upload is complete
-        attachment.upload_url = None
+        # Clear the upload descriptor since upload is complete
+        attachment.upload_descriptor = None
 
         # Save the updated attachment back to the store
         await data_store.save_attachment(attachment, {"user_id": DEFAULT_USER_ID})
@@ -637,19 +642,15 @@ async def preview_image(attachment_id: str):
         return JSONResponse(status_code=400, content={"error": "Invalid attachment ID."})
 
     try:
-        if not file_path.exists():
-            return JSONResponse(status_code=404, content={"error": "File not found"})
+        attachment = await data_store.load_attachment(attachment_id, {"user_id": DEFAULT_USER_ID})
+    except NotFoundError:
+        return JSONResponse(status_code=404, content={"error": "Attachment not found."})
 
-        # Determine media type from file extension or attachment metadata
-        # For simplicity, we'll try to load from the store
-        try:
-            attachment = await data_store.load_attachment(attachment_id, {"user_id": DEFAULT_USER_ID})
-            media_type = attachment.mime_type
-        except Exception:
-            # Default to binary if we can't determine
-            media_type = "application/octet-stream"
+    try:
+        if not file_path.exists():  # CodeQL [SM01305] Path is constrained by get_file_path.
+            return JSONResponse(status_code=404, content={"error": "File not found"})
 
-        return FileResponse(file_path, media_type=media_type)
+        return FileResponse(file_path, media_type=attachment.mime_type)  # CodeQL [SM01305] Path is constrained by get_file_path.  # fmt: skip
 
     except Exception as e:
         logger.error(f"Error serving preview for attachment {attachment_id}: {e}", exc_info=True)

python/samples/05-end-to-end/chatkit-integration/attachment_store.py

@@ -11,7 +11,13 @@
 from typing import TYPE_CHECKING, Any
 
 from chatkit.store import AttachmentStore
-from chatkit.types import Attachment, AttachmentCreateParams, FileAttachment, ImageAttachment
+from chatkit.types import (
+    Attachment,
+    AttachmentCreateParams,
+    AttachmentUploadDescriptor,
+    FileAttachment,
+    ImageAttachment,
+)
 from pydantic import AnyUrl
 
 if TYPE_CHECKING:
@@ -70,7 +76,7 @@ def get_file_path(self, attachment_id: str) -> Path:
         if not attachment_id or attachment_id in {".", ".."} or "/" in attachment_id or "\\" in attachment_id:
             raise ValueError(f"Invalid attachment ID: {attachment_id!r}")
 
-        file_path = (self.uploads_dir / attachment_id).resolve()
+        file_path = (self.uploads_dir / attachment_id).resolve()  # CodeQL [SM01305] Path containment is validated below.  # fmt: skip
         if not file_path.is_relative_to(self.uploads_dir) or file_path.parent != self.uploads_dir:
             raise ValueError(f"Invalid attachment ID: {attachment_id!r}")
         return file_path
@@ -90,8 +96,11 @@ async def create_attachment(self, input: AttachmentCreateParams, context: dict[s
         # Generate unique ID for this attachment
         attachment_id = self.generate_attachment_id(input.mime_type, context)
 
-        # Generate upload URL that points to our FastAPI upload endpoint
-        upload_url = f"{self.base_url}/upload/{attachment_id}"
+        # Generate upload instructions that point to our FastAPI upload endpoint
+        upload_descriptor = AttachmentUploadDescriptor(
+            url=AnyUrl(f"{self.base_url}/upload/{attachment_id}"),
+            method="POST",
+        )
 
         # Create appropriate attachment type based on MIME type
         if input.mime_type.startswith("image/"):
@@ -103,7 +112,7 @@ async def create_attachment(self, input: AttachmentCreateParams, context: dict[s
                 type="image",
                 mime_type=input.mime_type,
                 name=input.name,
-                upload_url=AnyUrl(upload_url),
+                upload_descriptor=upload_descriptor,
                 preview_url=AnyUrl(preview_url),
             )
         else:
@@ -113,7 +122,7 @@ async def create_attachment(self, input: AttachmentCreateParams, context: dict[s
                 type="file",
                 mime_type=input.mime_type,
                 name=input.name,
-                upload_url=AnyUrl(upload_url),
+                upload_descriptor=upload_descriptor,
             )
 
         # Save attachment metadata to data store so it's available during upload