feat(chat-streaming): added a stream option to workflow execute route, updated SDKs, updated docs (#1565)

* feat(chat-stream): updated workflow id execute route to support streaming via API * enable streaming via api * added only text stream option * cleanup deployed preview componnet * updated selectedOutputIds to selectedOutput * updated TS and Python SDKs with async, rate limits, usage, and streaming API routes * stream non-streaming blocks when streaming is specified * fix(chat-panel): add onBlockComplete handler to chat panel to stream back blocks as they complete * update docs * cleanup * ack PR comments * updated next config * removed getAssetUrl in favor of local assets * resolve merge conflicts * remove extra logic to create sensitive result * simplify internal auth * remove vercel blob from CSP + next config
2026-09-24 15:45:35 +08:00 · 2025-10-07 15:10:37 -07:00
parent a63a7b0262
commit 872e034312
73 changed files with 3662 additions and 1618 deletions
@@ -16,7 +16,7 @@ from simstudio import SimStudioClient

 # Initialize the client
 client = SimStudioClient(
-    api_key=os.getenv("SIMSTUDIO_API_KEY", "your-api-key-here"),
+    api_key=os.getenv("SIM_API_KEY", "your-api-key-here"),
    base_url="https://sim.ai"  # optional, defaults to https://sim.ai
 )

@@ -180,7 +180,7 @@ class SimStudioError(Exception):
 import os
 from simstudio import SimStudioClient

-client = SimStudioClient(api_key=os.getenv("SIMSTUDIO_API_KEY"))
+client = SimStudioClient(api_key=os.getenv("SIM_API_KEY"))

 def run_workflow():
    try:
@@ -216,7 +216,7 @@ run_workflow()
 from simstudio import SimStudioClient, SimStudioError
 import os

-client = SimStudioClient(api_key=os.getenv("SIMSTUDIO_API_KEY"))
+client = SimStudioClient(api_key=os.getenv("SIM_API_KEY"))

 def execute_with_error_handling():
    try:
@@ -246,7 +246,7 @@ from simstudio import SimStudioClient
 import os

 # Using context manager to automatically close the session
-with SimStudioClient(api_key=os.getenv("SIMSTUDIO_API_KEY")) as client:
+with SimStudioClient(api_key=os.getenv("SIM_API_KEY")) as client:
    result = client.execute_workflow("workflow-id")
    print("Result:", result)
 # Session is automatically closed here
@@ -260,8 +260,8 @@ from simstudio import SimStudioClient

 # Using environment variables
 client = SimStudioClient(
-    api_key=os.getenv("SIMSTUDIO_API_KEY"),
-    base_url=os.getenv("SIMSTUDIO_BASE_URL", "https://sim.ai")
+    api_key=os.getenv("SIM_API_KEY"),
+    base_url=os.getenv("SIM_BASE_URL", "https://sim.ai")
 )
 ```

@@ -271,7 +271,7 @@ client = SimStudioClient(
 from simstudio import SimStudioClient
 import os

-client = SimStudioClient(api_key=os.getenv("SIMSTUDIO_API_KEY"))
+client = SimStudioClient(api_key=os.getenv("SIM_API_KEY"))

 def execute_workflows_batch(workflow_data_pairs):
    """Execute multiple workflows with different input data."""
@@ -9,7 +9,7 @@ from simstudio import SimStudioClient, SimStudioError

 def basic_example():
    """Example 1: Basic workflow execution"""
-    client = SimStudioClient(api_key=os.getenv("SIMSTUDIO_API_KEY"))
+    client = SimStudioClient(api_key=os.getenv("SIM_API_KEY"))

    try:
        # Execute a workflow without input
@@ -31,7 +31,7 @@ def basic_example():

 def with_input_example():
    """Example 2: Workflow execution with input data"""
-    client = SimStudioClient(api_key=os.getenv("SIMSTUDIO_API_KEY"))
+    client = SimStudioClient(api_key=os.getenv("SIM_API_KEY"))

    try:
        result = client.execute_workflow(
@@ -66,7 +66,7 @@ def with_input_example():

 def status_example():
    """Example 3: Workflow validation and status checking"""
-    client = SimStudioClient(api_key=os.getenv("SIMSTUDIO_API_KEY"))
+    client = SimStudioClient(api_key=os.getenv("SIM_API_KEY"))

    try:
        # Check if workflow is ready
@@ -93,7 +93,7 @@ def status_example():

 def context_manager_example():
    """Example 4: Using context manager"""
-    with SimStudioClient(api_key=os.getenv("SIMSTUDIO_API_KEY")) as client:
+    with SimStudioClient(api_key=os.getenv("SIM_API_KEY")) as client:
        try:
            result = client.execute_workflow("your-workflow-id")
            print(f"Result: {result}")
@@ -104,7 +104,7 @@ def context_manager_example():

 def batch_execution_example():
    """Example 5: Batch workflow execution"""
-    client = SimStudioClient(api_key=os.getenv("SIMSTUDIO_API_KEY"))
+    client = SimStudioClient(api_key=os.getenv("SIM_API_KEY"))
    
    workflows = [
        ("workflow-1", {"type": "analysis", "data": "sample1"}),
@@ -155,13 +155,40 @@ def batch_execution_example():
    return results


+def streaming_example():
+    """Example 6: Workflow execution with streaming"""
+    client = SimStudioClient(api_key=os.getenv("SIM_API_KEY"))
+
+    try:
+        result = client.execute_workflow(
+            "your-workflow-id",
+            input_data={"message": "Count to five"},
+            stream=True,
+            selected_outputs=["agent1.content"],  # Use blockName.attribute format
+            timeout=60.0
+        )
+
+        if result.success:
+            print("✅ Workflow executed successfully!")
+            print(f"Output: {result.output}")
+            if result.metadata:
+                print(f"Duration: {result.metadata.get('duration')} ms")
+        else:
+            print(f"❌ Workflow failed: {result.error}")
+
+    except SimStudioError as error:
+        print(f"SDK Error: {error} (Code: {error.code})")
+    except Exception as error:
+        print(f"Unexpected error: {error}")
+
+
 def error_handling_example():
-    """Example 6: Comprehensive error handling"""
-    client = SimStudioClient(api_key=os.getenv("SIMSTUDIO_API_KEY"))
+    """Example 7: Comprehensive error handling"""
+    client = SimStudioClient(api_key=os.getenv("SIM_API_KEY"))

    try:
        result = client.execute_workflow("your-workflow-id")
-        
+
        if result.success:
            print("✅ Workflow executed successfully!")
            print(f"Output: {result.output}")
@@ -194,37 +221,41 @@ if __name__ == "__main__":
    print("🚀 Running Sim Python SDK Examples\n")
    
    # Check if API key is set
-    if not os.getenv("SIMSTUDIO_API_KEY"):
-        print("❌ Please set SIMSTUDIO_API_KEY environment variable")
+    if not os.getenv("SIM_API_KEY"):
+        print("❌ Please set SIM_API_KEY environment variable")
        exit(1)
    
    try:
        print("1️⃣ Basic Example:")
        basic_example()
        print("\n✅ Basic example completed\n")
-        
+
        print("2️⃣ Input Example:")
        with_input_example()
        print("\n✅ Input example completed\n")
-        
+
        print("3️⃣ Status Example:")
        status_example()
        print("\n✅ Status example completed\n")
-        
+
        print("4️⃣ Context Manager Example:")
        context_manager_example()
        print("\n✅ Context manager example completed\n")
-        
+
        print("5️⃣ Batch Execution Example:")
        batch_execution_example()
        print("\n✅ Batch execution example completed\n")
-        
-        print("6️⃣ Error Handling Example:")
+
+        print("6️⃣ Streaming Example:")
+        streaming_example()
+        print("\n✅ Streaming example completed\n")
+
+        print("7️⃣ Error Handling Example:")
        error_handling_example()
        print("\n✅ Error handling example completed\n")
-        
+
    except Exception as e:
        print(f"\n💥 Example failed: {e}")
        exit(1)
-    
+
    print("🎉 All examples completed successfully!") 
@@ -4,14 +4,24 @@ Sim SDK for Python
 Official Python SDK for Sim, allowing you to execute workflows programmatically.
 """

-from typing import Any, Dict, Optional
+from typing import Any, Dict, Optional, Union
 from dataclasses import dataclass
+import time
+import random

 import requests


 __version__ = "0.1.0"
-__all__ = ["SimStudioClient", "SimStudioError", "WorkflowExecutionResult", "WorkflowStatus"]
+__all__ = [
+    "SimStudioClient",
+    "SimStudioError",
+    "WorkflowExecutionResult",
+    "WorkflowStatus",
+    "AsyncExecutionResult",
+    "RateLimitInfo",
+    "UsageLimits",
+]


@dataclass
@@ -35,6 +45,42 @@ class WorkflowStatus:
    needs_redeployment: bool = False


+@dataclass
+class AsyncExecutionResult:
+    """Result of an async workflow execution."""
+    success: bool
+    task_id: str
+    status: str  # 'queued'
+    created_at: str
+    links: Dict[str, str]
+
+
+@dataclass
+class RateLimitInfo:
+    """Rate limit information from API response headers."""
+    limit: int
+    remaining: int
+    reset: int
+    retry_after: Optional[int] = None
+
+
+@dataclass
+class RateLimitStatus:
+    """Rate limit status for sync/async requests."""
+    is_limited: bool
+    limit: int
+    remaining: int
+    reset_at: str
+
+
+@dataclass
+class UsageLimits:
+    """Usage limits and quota information."""
+    success: bool
+    rate_limit: Dict[str, Any]
+    usage: Dict[str, Any]
+
+
 class SimStudioError(Exception):
    """Exception raised for Sim API errors."""
    
@@ -61,36 +107,69 @@ class SimStudioClient:
            'X-API-Key': self.api_key,
            'Content-Type': 'application/json',
        })
+        self._rate_limit_info: Optional[RateLimitInfo] = None
    
    def execute_workflow(
-        self, 
-        workflow_id: str, 
+        self,
+        workflow_id: str,
        input_data: Optional[Dict[str, Any]] = None,
-        timeout: float = 30.0
-    ) -> WorkflowExecutionResult:
+        timeout: float = 30.0,
+        stream: Optional[bool] = None,
+        selected_outputs: Optional[list] = None,
+        async_execution: Optional[bool] = None
+    ) -> Union[WorkflowExecutionResult, AsyncExecutionResult]:
        """
        Execute a workflow with optional input data.
-        
+        If async_execution is True, returns immediately with a task ID.
+
        Args:
            workflow_id: The ID of the workflow to execute
            input_data: Input data to pass to the workflow
            timeout: Timeout in seconds (default: 30.0)
-            
+            stream: Enable streaming responses (default: None)
+            selected_outputs: Block outputs to stream (e.g., ["agent1.content"])
+            async_execution: Execute asynchronously (default: None)
+
        Returns:
-            WorkflowExecutionResult object containing the execution result
-            
+            WorkflowExecutionResult or AsyncExecutionResult object
+
        Raises:
            SimStudioError: If the workflow execution fails
        """
        url = f"{self.base_url}/api/workflows/{workflow_id}/execute"
-        
+
+        # Build request body - spread input at root level, then add API control parameters
+        body = input_data.copy() if input_data is not None else {}
+        if stream is not None:
+            body['stream'] = stream
+        if selected_outputs is not None:
+            body['selectedOutputs'] = selected_outputs
+
+        # Build headers - async execution uses X-Execution-Mode header
+        headers = self._session.headers.copy()
+        if async_execution:
+            headers['X-Execution-Mode'] = 'async'
+
        try:
            response = self._session.post(
                url,
-                json=input_data or {},
+                json=body,
+                headers=headers,
                timeout=timeout
            )
-            
+
+            # Update rate limit info
+            self._update_rate_limit_info(response)
+
+            # Handle rate limiting
+            if response.status_code == 429:
+                retry_after = self._rate_limit_info.retry_after if self._rate_limit_info else 1000
+                raise SimStudioError(
+                    f'Rate limit exceeded. Retry after {retry_after}ms',
+                    'RATE_LIMIT_EXCEEDED',
+                    429
+                )
+
            if not response.ok:
                try:
                    error_data = response.json()
@@ -99,11 +178,21 @@ class SimStudioClient:
                except (ValueError, KeyError):
                    error_message = f'HTTP {response.status_code}: {response.reason}'
                    error_code = None
-                
+
                raise SimStudioError(error_message, error_code, response.status_code)
-            
+
            result_data = response.json()
-            
+
+            # Check if this is an async execution response (202 status)
+            if response.status_code == 202 and 'taskId' in result_data:
+                return AsyncExecutionResult(
+                    success=result_data.get('success', True),
+                    task_id=result_data['taskId'],
+                    status=result_data.get('status', 'queued'),
+                    created_at=result_data.get('createdAt', ''),
+                    links=result_data.get('links', {})
+                )
+
            return WorkflowExecutionResult(
                success=result_data['success'],
                output=result_data.get('output'),
@@ -113,7 +202,7 @@ class SimStudioClient:
                trace_spans=result_data.get('traceSpans'),
                total_duration=result_data.get('totalDuration')
            )
-            
+
        except requests.Timeout:
            raise SimStudioError(f'Workflow execution timed out after {timeout} seconds', 'TIMEOUT')
        except requests.RequestException as e:
@@ -180,28 +269,32 @@ class SimStudioClient:
        self,
        workflow_id: str,
        input_data: Optional[Dict[str, Any]] = None,
-        timeout: float = 30.0
+        timeout: float = 30.0,
+        stream: Optional[bool] = None,
+        selected_outputs: Optional[list] = None
    ) -> WorkflowExecutionResult:
        """
        Execute a workflow and poll for completion (useful for long-running workflows).
-        
+
        Note: Currently, the API is synchronous, so this method just calls execute_workflow.
        In the future, if async execution is added, this method can be enhanced.
-        
+
        Args:
            workflow_id: The ID of the workflow to execute
            input_data: Input data to pass to the workflow
            timeout: Timeout for the initial request in seconds
-            
+            stream: Enable streaming responses (default: None)
+            selected_outputs: Block outputs to stream (e.g., ["agent1.content"])
+
        Returns:
            WorkflowExecutionResult object containing the execution result
-            
+
        Raises:
            SimStudioError: If the workflow execution fails
        """
        # For now, the API is synchronous, so we just execute directly
        # In the future, if async execution is added, this method can be enhanced
-        return self.execute_workflow(workflow_id, input_data, timeout)
+        return self.execute_workflow(workflow_id, input_data, timeout, stream, selected_outputs)
    
    def set_api_key(self, api_key: str) -> None:
        """
@@ -225,11 +318,189 @@ class SimStudioClient:
    def close(self) -> None:
        """Close the underlying HTTP session."""
        self._session.close()
-    
+
+    def get_job_status(self, task_id: str) -> Dict[str, Any]:
+        """
+        Get the status of an async job.
+
+        Args:
+            task_id: The task ID returned from async execution
+
+        Returns:
+            Dictionary containing the job status
+
+        Raises:
+            SimStudioError: If getting the status fails
+        """
+        url = f"{self.base_url}/api/jobs/{task_id}"
+
+        try:
+            response = self._session.get(url)
+
+            self._update_rate_limit_info(response)
+
+            if not response.ok:
+                try:
+                    error_data = response.json()
+                    error_message = error_data.get('error', f'HTTP {response.status_code}: {response.reason}')
+                    error_code = error_data.get('code')
+                except (ValueError, KeyError):
+                    error_message = f'HTTP {response.status_code}: {response.reason}'
+                    error_code = None
+
+                raise SimStudioError(error_message, error_code, response.status_code)
+
+            return response.json()
+
+        except requests.RequestException as e:
+            raise SimStudioError(f'Failed to get job status: {str(e)}', 'STATUS_ERROR')
+
+    def execute_with_retry(
+        self,
+        workflow_id: str,
+        input_data: Optional[Dict[str, Any]] = None,
+        timeout: float = 30.0,
+        stream: Optional[bool] = None,
+        selected_outputs: Optional[list] = None,
+        async_execution: Optional[bool] = None,
+        max_retries: int = 3,
+        initial_delay: float = 1.0,
+        max_delay: float = 30.0,
+        backoff_multiplier: float = 2.0
+    ) -> Union[WorkflowExecutionResult, AsyncExecutionResult]:
+        """
+        Execute workflow with automatic retry on rate limit.
+
+        Args:
+            workflow_id: The ID of the workflow to execute
+            input_data: Input data to pass to the workflow
+            timeout: Timeout in seconds
+            stream: Enable streaming responses
+            selected_outputs: Block outputs to stream
+            async_execution: Execute asynchronously
+            max_retries: Maximum number of retries (default: 3)
+            initial_delay: Initial delay in seconds (default: 1.0)
+            max_delay: Maximum delay in seconds (default: 30.0)
+            backoff_multiplier: Backoff multiplier (default: 2.0)
+
+        Returns:
+            WorkflowExecutionResult or AsyncExecutionResult object
+
+        Raises:
+            SimStudioError: If max retries exceeded or other error occurs
+        """
+        last_error = None
+        delay = initial_delay
+
+        for attempt in range(max_retries + 1):
+            try:
+                return self.execute_workflow(
+                    workflow_id,
+                    input_data,
+                    timeout,
+                    stream,
+                    selected_outputs,
+                    async_execution
+                )
+            except SimStudioError as e:
+                if e.code != 'RATE_LIMIT_EXCEEDED':
+                    raise
+
+                last_error = e
+
+                # Don't retry after last attempt
+                if attempt == max_retries:
+                    break
+
+                # Use retry-after if provided, otherwise use exponential backoff
+                wait_time = (
+                    self._rate_limit_info.retry_after / 1000
+                    if self._rate_limit_info and self._rate_limit_info.retry_after
+                    else min(delay, max_delay)
+                )
+
+                # Add jitter (±25%)
+                jitter = wait_time * (0.75 + random.random() * 0.5)
+
+                time.sleep(jitter)
+
+                # Exponential backoff for next attempt
+                delay *= backoff_multiplier
+
+        raise last_error or SimStudioError('Max retries exceeded', 'MAX_RETRIES_EXCEEDED')
+
+    def get_rate_limit_info(self) -> Optional[RateLimitInfo]:
+        """
+        Get current rate limit information.
+
+        Returns:
+            RateLimitInfo object or None if no rate limit info available
+        """
+        return self._rate_limit_info
+
+    def _update_rate_limit_info(self, response: requests.Response) -> None:
+        """
+        Update rate limit info from response headers.
+
+        Args:
+            response: The response object to extract headers from
+        """
+        limit = response.headers.get('x-ratelimit-limit')
+        remaining = response.headers.get('x-ratelimit-remaining')
+        reset = response.headers.get('x-ratelimit-reset')
+        retry_after = response.headers.get('retry-after')
+
+        if limit or remaining or reset:
+            self._rate_limit_info = RateLimitInfo(
+                limit=int(limit) if limit else 0,
+                remaining=int(remaining) if remaining else 0,
+                reset=int(reset) if reset else 0,
+                retry_after=int(retry_after) * 1000 if retry_after else None
+            )
+
+    def get_usage_limits(self) -> UsageLimits:
+        """
+        Get current usage limits and quota information.
+
+        Returns:
+            UsageLimits object containing usage and quota data
+
+        Raises:
+            SimStudioError: If getting usage limits fails
+        """
+        url = f"{self.base_url}/api/users/me/usage-limits"
+
+        try:
+            response = self._session.get(url)
+
+            self._update_rate_limit_info(response)
+
+            if not response.ok:
+                try:
+                    error_data = response.json()
+                    error_message = error_data.get('error', f'HTTP {response.status_code}: {response.reason}')
+                    error_code = error_data.get('code')
+                except (ValueError, KeyError):
+                    error_message = f'HTTP {response.status_code}: {response.reason}'
+                    error_code = None
+
+                raise SimStudioError(error_message, error_code, response.status_code)
+
+            data = response.json()
+
+            return UsageLimits(
+                success=data.get('success', True),
+                rate_limit=data.get('rateLimit', {}),
+                usage=data.get('usage', {})
+            )
+
+        except requests.RequestException as e:
+            raise SimStudioError(f'Failed to get usage limits: {str(e)}', 'USAGE_ERROR')
+
    def __enter__(self):
        """Context manager entry."""
        return self
-    
+
    def __exit__(self, exc_type, exc_val, exc_tb):
        """Context manager exit."""
        self.close()
@@ -94,4 +94,371 @@ def test_context_manager(mock_close):
    with SimStudioClient(api_key="test-api-key") as client:
        assert client.api_key == "test-api-key"
    # Should close without error
-    mock_close.assert_called_once() 
+    mock_close.assert_called_once()
+
+
+# Tests for async execution
+@patch('simstudio.requests.Session.post')
+def test_async_execution_returns_task_id(mock_post):
+    """Test async execution returns AsyncExecutionResult."""
+    mock_response = Mock()
+    mock_response.ok = True
+    mock_response.status_code = 202
+    mock_response.json.return_value = {
+        "success": True,
+        "taskId": "task-123",
+        "status": "queued",
+        "createdAt": "2024-01-01T00:00:00Z",
+        "links": {"status": "/api/jobs/task-123"}
+    }
+    mock_response.headers.get.return_value = None
+    mock_post.return_value = mock_response
+
+    client = SimStudioClient(api_key="test-api-key")
+    result = client.execute_workflow(
+        "workflow-id",
+        input_data={"message": "Hello"},
+        async_execution=True
+    )
+
+    assert result.success is True
+    assert result.task_id == "task-123"
+    assert result.status == "queued"
+    assert result.links["status"] == "/api/jobs/task-123"
+
+    # Verify X-Execution-Mode header was set
+    call_args = mock_post.call_args
+    assert call_args[1]["headers"]["X-Execution-Mode"] == "async"
+
+
+@patch('simstudio.requests.Session.post')
+def test_sync_execution_returns_result(mock_post):
+    """Test sync execution returns WorkflowExecutionResult."""
+    mock_response = Mock()
+    mock_response.ok = True
+    mock_response.status_code = 200
+    mock_response.json.return_value = {
+        "success": True,
+        "output": {"result": "completed"},
+        "logs": []
+    }
+    mock_response.headers.get.return_value = None
+    mock_post.return_value = mock_response
+
+    client = SimStudioClient(api_key="test-api-key")
+    result = client.execute_workflow(
+        "workflow-id",
+        input_data={"message": "Hello"},
+        async_execution=False
+    )
+
+    assert result.success is True
+    assert result.output == {"result": "completed"}
+    assert not hasattr(result, 'task_id')
+
+
+@patch('simstudio.requests.Session.post')
+def test_async_header_not_set_when_false(mock_post):
+    """Test X-Execution-Mode header is not set when async_execution is None."""
+    mock_response = Mock()
+    mock_response.ok = True
+    mock_response.status_code = 200
+    mock_response.json.return_value = {"success": True, "output": {}}
+    mock_response.headers.get.return_value = None
+    mock_post.return_value = mock_response
+
+    client = SimStudioClient(api_key="test-api-key")
+    client.execute_workflow("workflow-id", input_data={"message": "Hello"})
+
+    call_args = mock_post.call_args
+    assert "X-Execution-Mode" not in call_args[1]["headers"]
+
+
+# Tests for job status
+@patch('simstudio.requests.Session.get')
+def test_get_job_status_success(mock_get):
+    """Test getting job status."""
+    mock_response = Mock()
+    mock_response.ok = True
+    mock_response.json.return_value = {
+        "success": True,
+        "taskId": "task-123",
+        "status": "completed",
+        "metadata": {
+            "startedAt": "2024-01-01T00:00:00Z",
+            "completedAt": "2024-01-01T00:01:00Z",
+            "duration": 60000
+        },
+        "output": {"result": "done"}
+    }
+    mock_response.headers.get.return_value = None
+    mock_get.return_value = mock_response
+
+    client = SimStudioClient(api_key="test-api-key", base_url="https://test.sim.ai")
+    result = client.get_job_status("task-123")
+
+    assert result["taskId"] == "task-123"
+    assert result["status"] == "completed"
+    assert result["output"]["result"] == "done"
+    mock_get.assert_called_once_with("https://test.sim.ai/api/jobs/task-123")
+
+
+@patch('simstudio.requests.Session.get')
+def test_get_job_status_not_found(mock_get):
+    """Test job not found error."""
+    mock_response = Mock()
+    mock_response.ok = False
+    mock_response.status_code = 404
+    mock_response.reason = "Not Found"
+    mock_response.json.return_value = {
+        "error": "Job not found",
+        "code": "JOB_NOT_FOUND"
+    }
+    mock_response.headers.get.return_value = None
+    mock_get.return_value = mock_response
+
+    client = SimStudioClient(api_key="test-api-key")
+
+    with pytest.raises(SimStudioError) as exc_info:
+        client.get_job_status("invalid-task")
+    assert "Job not found" in str(exc_info.value)
+
+
+# Tests for retry with rate limiting
+@patch('simstudio.requests.Session.post')
+@patch('simstudio.time.sleep')
+def test_execute_with_retry_success_first_attempt(mock_sleep, mock_post):
+    """Test retry succeeds on first attempt."""
+    mock_response = Mock()
+    mock_response.ok = True
+    mock_response.status_code = 200
+    mock_response.json.return_value = {
+        "success": True,
+        "output": {"result": "success"}
+    }
+    mock_response.headers.get.return_value = None
+    mock_post.return_value = mock_response
+
+    client = SimStudioClient(api_key="test-api-key")
+    result = client.execute_with_retry("workflow-id", input_data={"message": "test"})
+
+    assert result.success is True
+    assert mock_post.call_count == 1
+    assert mock_sleep.call_count == 0
+
+
+@patch('simstudio.requests.Session.post')
+@patch('simstudio.time.sleep')
+def test_execute_with_retry_retries_on_rate_limit(mock_sleep, mock_post):
+    """Test retry retries on rate limit error."""
+    rate_limit_response = Mock()
+    rate_limit_response.ok = False
+    rate_limit_response.status_code = 429
+    rate_limit_response.json.return_value = {
+        "error": "Rate limit exceeded",
+        "code": "RATE_LIMIT_EXCEEDED"
+    }
+    import time
+    rate_limit_response.headers.get.side_effect = lambda h: {
+        'retry-after': '1',
+        'x-ratelimit-limit': '100',
+        'x-ratelimit-remaining': '0',
+        'x-ratelimit-reset': str(int(time.time()) + 60)
+    }.get(h)
+
+    success_response = Mock()
+    success_response.ok = True
+    success_response.status_code = 200
+    success_response.json.return_value = {
+        "success": True,
+        "output": {"result": "success"}
+    }
+    success_response.headers.get.return_value = None
+
+    mock_post.side_effect = [rate_limit_response, success_response]
+
+    client = SimStudioClient(api_key="test-api-key")
+    result = client.execute_with_retry(
+        "workflow-id",
+        input_data={"message": "test"},
+        max_retries=3,
+        initial_delay=0.01
+    )
+
+    assert result.success is True
+    assert mock_post.call_count == 2
+    assert mock_sleep.call_count == 1
+
+
+@patch('simstudio.requests.Session.post')
+@patch('simstudio.time.sleep')
+def test_execute_with_retry_max_retries_exceeded(mock_sleep, mock_post):
+    """Test retry throws after max retries."""
+    mock_response = Mock()
+    mock_response.ok = False
+    mock_response.status_code = 429
+    mock_response.json.return_value = {
+        "error": "Rate limit exceeded",
+        "code": "RATE_LIMIT_EXCEEDED"
+    }
+    mock_response.headers.get.side_effect = lambda h: '1' if h == 'retry-after' else None
+    mock_post.return_value = mock_response
+
+    client = SimStudioClient(api_key="test-api-key")
+
+    with pytest.raises(SimStudioError) as exc_info:
+        client.execute_with_retry(
+            "workflow-id",
+            input_data={"message": "test"},
+            max_retries=2,
+            initial_delay=0.01
+        )
+
+    assert "Rate limit exceeded" in str(exc_info.value)
+    assert mock_post.call_count == 3  # Initial + 2 retries
+
+
+@patch('simstudio.requests.Session.post')
+def test_execute_with_retry_no_retry_on_other_errors(mock_post):
+    """Test retry does not retry on non-rate-limit errors."""
+    mock_response = Mock()
+    mock_response.ok = False
+    mock_response.status_code = 500
+    mock_response.reason = "Internal Server Error"
+    mock_response.json.return_value = {
+        "error": "Server error",
+        "code": "INTERNAL_ERROR"
+    }
+    mock_response.headers.get.return_value = None
+    mock_post.return_value = mock_response
+
+    client = SimStudioClient(api_key="test-api-key")
+
+    with pytest.raises(SimStudioError) as exc_info:
+        client.execute_with_retry("workflow-id", input_data={"message": "test"})
+
+    assert "Server error" in str(exc_info.value)
+    assert mock_post.call_count == 1  # No retries
+
+
+# Tests for rate limit info
+def test_get_rate_limit_info_returns_none_initially():
+    """Test rate limit info is None before any API calls."""
+    client = SimStudioClient(api_key="test-api-key")
+    info = client.get_rate_limit_info()
+    assert info is None
+
+
+@patch('simstudio.requests.Session.post')
+def test_get_rate_limit_info_after_api_call(mock_post):
+    """Test rate limit info is populated after API call."""
+    mock_response = Mock()
+    mock_response.ok = True
+    mock_response.status_code = 200
+    mock_response.json.return_value = {"success": True, "output": {}}
+    mock_response.headers.get.side_effect = lambda h: {
+        'x-ratelimit-limit': '100',
+        'x-ratelimit-remaining': '95',
+        'x-ratelimit-reset': '1704067200'
+    }.get(h)
+    mock_post.return_value = mock_response
+
+    client = SimStudioClient(api_key="test-api-key")
+    client.execute_workflow("workflow-id", input_data={})
+
+    info = client.get_rate_limit_info()
+    assert info is not None
+    assert info.limit == 100
+    assert info.remaining == 95
+    assert info.reset == 1704067200
+
+
+# Tests for usage limits
+@patch('simstudio.requests.Session.get')
+def test_get_usage_limits_success(mock_get):
+    """Test getting usage limits."""
+    mock_response = Mock()
+    mock_response.ok = True
+    mock_response.json.return_value = {
+        "success": True,
+        "rateLimit": {
+            "sync": {
+                "isLimited": False,
+                "limit": 100,
+                "remaining": 95,
+                "resetAt": "2024-01-01T01:00:00Z"
+            },
+            "async": {
+                "isLimited": False,
+                "limit": 50,
+                "remaining": 48,
+                "resetAt": "2024-01-01T01:00:00Z"
+            },
+            "authType": "api"
+        },
+        "usage": {
+            "currentPeriodCost": 1.23,
+            "limit": 100.0,
+            "plan": "pro"
+        }
+    }
+    mock_response.headers.get.return_value = None
+    mock_get.return_value = mock_response
+
+    client = SimStudioClient(api_key="test-api-key", base_url="https://test.sim.ai")
+    result = client.get_usage_limits()
+
+    assert result.success is True
+    assert result.rate_limit["sync"]["limit"] == 100
+    assert result.rate_limit["async"]["limit"] == 50
+    assert result.usage["currentPeriodCost"] == 1.23
+    assert result.usage["plan"] == "pro"
+    mock_get.assert_called_once_with("https://test.sim.ai/api/users/me/usage-limits")
+
+
+@patch('simstudio.requests.Session.get')
+def test_get_usage_limits_unauthorized(mock_get):
+    """Test usage limits with invalid API key."""
+    mock_response = Mock()
+    mock_response.ok = False
+    mock_response.status_code = 401
+    mock_response.reason = "Unauthorized"
+    mock_response.json.return_value = {
+        "error": "Invalid API key",
+        "code": "UNAUTHORIZED"
+    }
+    mock_response.headers.get.return_value = None
+    mock_get.return_value = mock_response
+
+    client = SimStudioClient(api_key="invalid-key")
+
+    with pytest.raises(SimStudioError) as exc_info:
+        client.get_usage_limits()
+    assert "Invalid API key" in str(exc_info.value)
+
+
+# Tests for streaming with selectedOutputs
+@patch('simstudio.requests.Session.post')
+def test_execute_workflow_with_stream_and_selected_outputs(mock_post):
+    """Test execution with stream and selectedOutputs parameters."""
+    mock_response = Mock()
+    mock_response.ok = True
+    mock_response.status_code = 200
+    mock_response.json.return_value = {"success": True, "output": {}}
+    mock_response.headers.get.return_value = None
+    mock_post.return_value = mock_response
+
+    client = SimStudioClient(api_key="test-api-key")
+    client.execute_workflow(
+        "workflow-id",
+        input_data={"message": "test"},
+        stream=True,
+        selected_outputs=["agent1.content", "agent2.content"]
+    )
+
+    call_args = mock_post.call_args
+    request_body = call_args[1]["json"]
+
+    assert request_body["message"] == "test"
+    assert request_body["stream"] is True
+    assert request_body["selectedOutputs"] == ["agent1.content", "agent2.content"]