From 9efe1df9172c4c657f7fd27b9198d265852cc638 Mon Sep 17 00:00:00 2001
From: Joey French <joey@harbinger.finance>
Date: Thu, 27 Nov 2025 16:16:34 -0600
Subject: [PATCH] feat: add S3 endpoint configuration to RoboSystemsExtension

- Introduced `s3_endpoint_url` to `RoboSystemsExtensionConfig` for optional S3 endpoint overrides, facilitating integration with services like LocalStack.
- Updated `FileClient` to utilize the new `s3_endpoint_url` for overriding upload URLs, enhancing flexibility in file upload operations.

These changes improve the configurability of S3 interactions within the client.
---
 robosystems_client/extensions/extensions.py  |  2 ++
 robosystems_client/extensions/file_client.py | 24 ++++++++++++++++----
 2 files changed, 22 insertions(+), 4 deletions(-)

diff --git a/robosystems_client/extensions/extensions.py b/robosystems_client/extensions/extensions.py
index f1d031c..ddf53d9 100644
--- a/robosystems_client/extensions/extensions.py
+++ b/robosystems_client/extensions/extensions.py
@@ -25,6 +25,7 @@ class RoboSystemsExtensionConfig:
   max_retries: int = 5
   retry_delay: int = 1000
   timeout: int = 30
+  s3_endpoint_url: Optional[str] = None  # Override S3 endpoint (e.g., for LocalStack)
 
 
 class RoboSystemsExtensions:
@@ -41,6 +42,7 @@ def __init__(self, config: RoboSystemsExtensionConfig = None):
       "max_retries": config.max_retries,
       "retry_delay": config.retry_delay,
       "timeout": config.timeout,
+      "s3_endpoint_url": config.s3_endpoint_url,
     }
 
     # Extract token from headers if it was set by auth classes
diff --git a/robosystems_client/extensions/file_client.py b/robosystems_client/extensions/file_client.py
index b5a40a7..ace3706 100644
--- a/robosystems_client/extensions/file_client.py
+++ b/robosystems_client/extensions/file_client.py
@@ -37,7 +37,6 @@ class FileUploadOptions:
   """Options for file upload operations"""
 
   on_progress: Optional[Callable[[str], None]] = None
-  fix_localstack_url: bool = True
   ingest_to_graph: bool = False
 
 
@@ -78,6 +77,9 @@ def __init__(self, config: Dict[str, Any]):
     self.base_url = config["base_url"]
     self.headers = config.get("headers", {})
     self.token = config.get("token")
+    self.s3_endpoint_url = config.get(
+      "s3_endpoint_url"
+    )  # Optional S3 endpoint override
     self._http_client = httpx.Client(timeout=120.0)
 
   def upload(
@@ -171,9 +173,23 @@ def upload(
       upload_url = upload_data.upload_url
       file_id = upload_data.file_id
 
-      # Fix LocalStack URL if needed
-      if options.fix_localstack_url and "localstack:4566" in upload_url:
-        upload_url = upload_url.replace("localstack:4566", "localhost:4566")
+      # Override S3 endpoint if configured (e.g., for LocalStack)
+      if self.s3_endpoint_url:
+        from urllib.parse import urlparse, urlunparse
+
+        parsed_url = urlparse(upload_url)
+        override_parsed = urlparse(self.s3_endpoint_url)
+        # Replace scheme, host, and port with the override endpoint
+        upload_url = urlunparse(
+          (
+            override_parsed.scheme or parsed_url.scheme,
+            override_parsed.netloc,
+            parsed_url.path,
+            parsed_url.params,
+            parsed_url.query,
+            parsed_url.fragment,
+          )
+        )
 
       # Step 2: Upload file to S3
       if options.on_progress: