fix UTs

Author: liqul

pytest.ini

@@ -1,5 +1,9 @@
 [pytest]
+pythonpath = .
+asyncio_mode = auto
 markers =
     app_config: mark a test that requires the app config
+    integration: mark a test as an integration test (may start real services)
+    asyncio: mark a test as async
 testpaths =
     tests

taskweaver/ces/AGENTS.md

@@ -171,6 +171,7 @@ Connect to pre-started server. API key required.
 | POST | `/api/v1/sessions/{id}/execute` | Execute code |
 | GET | `/api/v1/sessions/{id}/stream/{exec_id}` | SSE stream |
 | POST | `/api/v1/sessions/{id}/variables` | Update variables |
+| POST | `/api/v1/sessions/{id}/files` | Upload file to session cwd |
 | GET | `/api/v1/sessions/{id}/artifacts/{file}` | Download artifact |
 
 ## Usage
@@ -236,6 +237,140 @@ with ExecutionClient(
 4. **Streaming**: SSE events for stdout/stderr during execution
 5. **Session Cleanup**: `DELETE /sessions/{id}` → Environment.stop_session()
 
+## File Upload Flow
+
+File upload enables the `/load` CLI command to transfer files from the client machine to the execution server's working directory. This is essential when the execution server runs in a container or on a remote machine where the client's local filesystem is not accessible.
+
+### Architecture
+
+```
+┌─────────────────────────────────────────────────────────────────────────────┐
+│                          TASKWEAVER CLIENT                                  │
+│                                                                             │
+│  ┌─────────────┐    ┌─────────────┐    ┌───────────────────────────────┐   │
+│  │   Session   │───▶│ _upload_file│───▶│ ExecutionServiceClient        │   │
+│  │  /load cmd  │    │   (lazy)    │    │   .upload_file()              │   │
+│  └─────────────┘    └─────────────┘    └───────────────┬───────────────┘   │
+│                                                        │                    │
+│                                                        ▼                    │
+│                                        ┌───────────────────────────────┐   │
+│                                        │ ExecutionClient.upload_file() │   │
+│                                        │  - Read file content          │   │
+│                                        │  - Base64 encode              │   │
+│                                        │  - HTTP POST                  │   │
+│                                        └───────────────┬───────────────┘   │
+└────────────────────────────────────────────────────────┼────────────────────┘
+                                                         │
+                                                         │ POST /api/v1/sessions/{id}/files
+                                                         │ {filename, content (base64), encoding}
+                                                         ▼
+┌─────────────────────────────────────────────────────────────────────────────┐
+│                          EXECUTION SERVER                                   │
+│                                                                             │
+│  ┌───────────────────────────────────────────────────────────────────────┐  │
+│  │                    routes.upload_file()                               │  │
+│  │  - Validate session exists                                            │  │
+│  │  - Base64 decode content                                              │  │
+│  │  - Call session_manager.upload_file()                                 │  │
+│  └───────────────────────────────────────────────────────────────────────┘  │
+│                                    │                                        │
+│                                    ▼                                        │
+│  ┌───────────────────────────────────────────────────────────────────────┐  │
+│  │              ServerSessionManager.upload_file()                       │  │
+│  │  - Sanitize filename (prevent path traversal)                         │  │
+│  │  - Write to {session.cwd}/{filename}                                  │  │
+│  │  - Return full path                                                   │  │
+│  └───────────────────────────────────────────────────────────────────────┘  │
+│                                    │                                        │
+│                                    ▼                                        │
+│  ┌───────────────────────────────────────────────────────────────────────┐  │
+│  │                    Session Working Directory                          │  │
+│  │                    /workspace/{session_id}/cwd/                       │  │
+│  │                         └── uploaded_file.csv                         │  │
+│  └───────────────────────────────────────────────────────────────────────┘  │
+└─────────────────────────────────────────────────────────────────────────────┘
+```
+
+### Request/Response Models
+
+```python
+# Request (server/models.py)
+class UploadFileRequest(BaseModel):
+    filename: str       # Target filename (basename extracted, path traversal prevented)
+    content: str        # File content (base64 encoded for binary)
+    encoding: Literal["base64", "text"] = "base64"
+
+# Response (server/models.py)
+class UploadFileResponse(BaseModel):
+    filename: str       # Uploaded filename
+    status: Literal["uploaded"] = "uploaded"
+    path: str           # Full path where file was saved on server
+```
+
+### Client Usage
+
+```python
+from taskweaver.ces.client import ExecutionClient
+
+with ExecutionClient(session_id="my-session", server_url="http://localhost:8000") as client:
+    client.start()
+    
+    # Upload a file
+    with open("/local/path/data.csv", "rb") as f:
+        content = f.read()
+    saved_path = client.upload_file("data.csv", content)
+    
+    # Now the file is available in the session's cwd
+    result = client.execute_code("exec-1", "import pandas as pd; df = pd.read_csv('data.csv')")
+```
+
+### Session Integration
+
+The `Session` class uses a lazily-initialized upload client:
+
+```python
+# In taskweaver/session/session.py
+class Session:
+    def _get_upload_client(self):
+        """Lazy client creation - only created when first upload occurs."""
+        if not hasattr(self, "_upload_client"):
+            self._upload_client = self.exec_mgr.get_session_client(
+                self.session_id,
+                session_dir=self.workspace,
+                cwd=self.execution_cwd,
+            )
+            self._upload_client_started = False
+        
+        if not self._upload_client_started:
+            self._upload_client.start()
+            self._upload_client_started = True
+        
+        return self._upload_client
+
+    def _upload_file(self, name: str, path: str = None, content: bytes = None) -> str:
+        """Upload file to execution server."""
+        target_name = os.path.basename(name)
+        
+        if path is not None:
+            with open(path, "rb") as f:
+                file_content = f.read()
+        elif content is not None:
+            file_content = content
+        else:
+            raise ValueError("path or content must be provided")
+        
+        client = self._get_upload_client()
+        client.upload_file(target_name, file_content)
+        return target_name
+```
+
+### Security Considerations
+
+1. **Path Traversal Prevention**: Server sanitizes filename using `os.path.basename()` to prevent `../../etc/passwd` attacks
+2. **Session Isolation**: Files are written only to the session's own cwd directory
+3. **API Key Authentication**: Upload endpoint respects the same API key auth as other endpoints
+4. **Size Limits**: Large files should be chunked or streamed (not yet implemented)
+
 ## Custom Kernel Magics (kernel/ext.py)
 
 ```python
@@ -271,6 +406,11 @@ Unit tests in `tests/unit_tests/ces/`:
 | `test_server_launcher.py` | ServerLauncher (mocked subprocess/docker) |
 | `test_execution_service.py` | ExecutionServiceProvider |
 
+**TODO**: Add tests for file upload functionality:
+- `ExecutionClient.upload_file()` - mock HTTP POST, verify base64 encoding
+- `ServerSessionManager.upload_file()` - verify file written, path traversal blocked
+- `routes.upload_file()` - integration test with mocked session manager
+
 Run tests:
 ```bash
 pytest tests/unit_tests/ces/ -v

taskweaver/ces/client/execution_client.py

@@ -431,6 +431,37 @@ def download_artifact(self, filename: str) -> bytes:
             )
         return response.content
 
+    def upload_file(
+        self,
+        filename: str,
+        content: bytes,
+    ) -> str:
+        """Upload a file to the session's working directory.
+
+        Args:
+            filename: Target filename.
+            content: File content as bytes.
+
+        Returns:
+            Path where the file was saved on the server.
+
+        Raises:
+            ExecutionClientError: If upload fails.
+        """
+        import base64
+
+        response = self._client.post(
+            f"/api/v1/sessions/{self.session_id}/files",
+            json={
+                "filename": filename,
+                "content": base64.b64encode(content).decode("ascii"),
+                "encoding": "base64",
+            },
+        )
+        result = self._handle_response(response)
+        logger.info(f"Uploaded file {filename} to session {self.session_id}")
+        return result.get("path", "")
+
     def close(self) -> None:
         """Close the HTTP client and release resources."""
         self._client.close()

taskweaver/ces/common.py

@@ -111,6 +111,14 @@ def execute_code(
     ) -> ExecutionResult:
         ...
 
+    @abstractmethod
+    def upload_file(
+        self,
+        filename: str,
+        content: bytes,
+    ) -> str:
+        ...
+
 
 KernelModeType = Literal["local", "container"]
 

taskweaver/ces/manager/defer.py

@@ -91,6 +91,13 @@ def execute_code(
     ) -> ExecutionResult:
         return self._get_proxy_client().execute_code(exec_id, code, on_output=on_output)
 
+    def upload_file(
+        self,
+        filename: str,
+        content: bytes,
+    ) -> str:
+        return self._get_proxy_client().upload_file(filename, content)
+
     def _get_proxy_client(self) -> Client:
         return self._init_deferred_var()()
 

taskweaver/ces/manager/execution_service.py

@@ -107,6 +107,16 @@ def execute_code(
             raise RuntimeError("Client not started")
         return self._client.execute_code(exec_id, code, on_output=on_output)
 
+    def upload_file(
+        self,
+        filename: str,
+        content: bytes,
+    ) -> str:
+        """Upload a file to the session's working directory."""
+        if self._client is None:
+            raise RuntimeError("Client not started")
+        return self._client.upload_file(filename, content)
+
 
 class ExecutionServiceProvider(Manager):
     """Manager implementation that uses the HTTP execution server.

taskweaver/ces/manager/sub_proc.py

@@ -60,6 +60,25 @@ def execute_code(
             on_output=on_output,
         )
 
+    def upload_file(
+        self,
+        filename: str,
+        content: bytes,
+    ) -> str:
+        """Upload a file to the session's working directory.
+
+        For subprocess mode, this writes directly to the local filesystem
+        since the subprocess shares the same filesystem as the caller.
+        """
+        # Sanitize filename to prevent path traversal
+        safe_filename = os.path.basename(filename)
+        file_path = os.path.join(self.cwd, safe_filename)
+
+        with open(file_path, "wb") as f:
+            f.write(content)
+
+        return file_path
+
 
 class SubProcessManager(Manager):
     def __init__(

taskweaver/ces/runtime/context.py

@@ -157,15 +157,22 @@ def get_session_var(
 
     def extract_visible_variables(self, local_ns: Dict[str, Any]) -> List[Tuple[str, str]]:
         ignore_names = {
+            # IPython/Jupyter internals
             "__builtins__",
             "In",
             "Out",
             "get_ipython",
             "exit",
             "quit",
+            # Common library aliases
             "pd",
             "np",
             "plt",
+            # REPL/kernel internals
+            "original_ps1",
+            "is_wsl",
+            "PS1",
+            "REPLHooks",
         }
 
         visible: List[Tuple[str, str]] = []
@@ -194,7 +201,12 @@ def extract_visible_variables(self, local_ns: Dict[str, Any]) -> List[Tuple[str,
                 continue
 
             try:
-                rendered = repr(value)
+                # Use str() for strings to avoid extra quotes from repr()
+                # e.g., repr("hello") returns "'hello'" but str("hello") returns "hello"
+                if isinstance(value, str):
+                    rendered = value
+                else:
+                    rendered = repr(value)
             except Exception:
                 rendered = "<unrepresentable>"
 

taskweaver/ces/server/models.py

@@ -39,6 +39,14 @@ class UpdateVariablesRequest(BaseModel):
     variables: Dict[str, str] = Field(..., description="Session variables to update")
 
 
+class UploadFileRequest(BaseModel):
+    """Request to upload a file to a session's working directory."""
+
+    filename: str = Field(..., description="Target filename in the session's cwd")
+    content: str = Field(..., description="File content (base64 encoded for binary, plain for text)")
+    encoding: Literal["base64", "text"] = Field("base64", description="Content encoding")
+
+
 # =============================================================================
 # Response Models
 # =============================================================================
@@ -128,6 +136,14 @@ class UpdateVariablesResponse(BaseModel):
     variables: Dict[str, str] = Field(..., description="Updated variables")
 
 
+class UploadFileResponse(BaseModel):
+    """Response after uploading a file."""
+
+    filename: str = Field(..., description="Uploaded filename")
+    status: Literal["uploaded"] = "uploaded"
+    path: str = Field(..., description="Path where file was saved")
+
+
 class ErrorResponse(BaseModel):
     """Standard error response."""
 

taskweaver/ces/server/routes.py

@@ -26,6 +26,8 @@
     StopSessionResponse,
     UpdateVariablesRequest,
     UpdateVariablesResponse,
+    UploadFileRequest,
+    UploadFileResponse,
     execution_result_to_response,
 )
 from taskweaver.ces.server.session_manager import ServerSessionManager
@@ -439,6 +441,39 @@ async def update_variables(
         raise HTTPException(status_code=500, detail=str(e))
 
 
+@router.post(
+    "/sessions/{session_id}/files",
+    response_model=UploadFileResponse,
+    dependencies=[Depends(verify_api_key)],
+)
+async def upload_file(
+    session_id: str,
+    request: UploadFileRequest,
+    session_manager: ServerSessionManager = Depends(get_session_manager),
+) -> UploadFileResponse:
+    """Upload a file to a session's working directory."""
+    if not session_manager.session_exists(session_id):
+        raise HTTPException(status_code=404, detail=f"Session {session_id} not found")
+
+    try:
+        import base64
+
+        if request.encoding == "base64":
+            content = base64.b64decode(request.content)
+        else:
+            content = request.content.encode("utf-8")
+
+        file_path = session_manager.upload_file(session_id, request.filename, content)
+        return UploadFileResponse(
+            filename=request.filename,
+            status="uploaded",
+            path=file_path,
+        )
+    except Exception as e:
+        logger.error(f"Failed to upload file {request.filename} to session {session_id}: {e}")
+        raise HTTPException(status_code=500, detail=str(e))
+
+
 # =============================================================================
 # Artifacts
 # =============================================================================