diff --git a/packages/prime-sandboxes/src/prime_sandboxes/models.py b/packages/prime-sandboxes/src/prime_sandboxes/models.py index acdec9e2..26d1b422 100644 --- a/packages/prime-sandboxes/src/prime_sandboxes/models.py +++ b/packages/prime-sandboxes/src/prime_sandboxes/models.py @@ -175,6 +175,7 @@ class CommandRequest(BaseModel): command: str working_dir: Optional[str] = None env: Optional[Dict[str, str]] = None + user: Optional[str] = None class CommandResponse(BaseModel): diff --git a/packages/prime-sandboxes/src/prime_sandboxes/sandbox.py b/packages/prime-sandboxes/src/prime_sandboxes/sandbox.py index b36092fb..e64fae6b 100644 --- a/packages/prime-sandboxes/src/prime_sandboxes/sandbox.py +++ b/packages/prime-sandboxes/src/prime_sandboxes/sandbox.py @@ -721,11 +721,17 @@ def execute_command( working_dir: Optional[str] = None, env: Optional[Dict[str, str]] = None, timeout: Optional[int] = None, + user: Optional[str] = None, ) -> CommandResponse: """Execute command directly via gateway.""" auth = self._auth_cache.get_or_refresh(sandbox_id) if self._auth_cache.is_vm(sandbox_id): + if user is not None: + raise ValueError( + "The 'user' parameter is only supported for container sandboxes, " + "not VM sandboxes." + ) return self._execute_command_connect_rpc( sandbox_id=sandbox_id, command=command, @@ -742,6 +748,7 @@ def execute_command( working_dir=working_dir, env=env, timeout=timeout, + user=user, ) def _execute_command_connect_rpc( @@ -822,6 +829,7 @@ def _execute_command_rest( working_dir: Optional[str] = None, env: Optional[Dict[str, str]] = None, timeout: Optional[int] = None, + user: Optional[str] = None, ) -> CommandResponse: gateway_url = auth["gateway_url"].rstrip("/") url = f"{gateway_url}/{auth['user_ns']}/{auth['job_id']}/exec" @@ -835,6 +843,8 @@ def _execute_command_rest( "sandbox_id": sandbox_id, "timeout": effective_timeout, } + if user is not None: + payload["user"] = user for attempt in range(MAX_409_RETRIES): try: @@ -900,6 +910,7 @@ def start_background_job( command: str, working_dir: Optional[str] = None, env: Optional[Dict[str, str]] = None, + user: Optional[str] = None, ) -> BackgroundJob: """Start a long-running command in the background. @@ -911,6 +922,8 @@ def start_background_job( command: Command to execute working_dir: Working directory for command execution env: Environment variables + user: Run the job as this user, like ``docker exec -u`` (username or + numeric UID, optionally USER:GROUP). Container sandboxes only. Returns: BackgroundJob with job_id and file paths for polling @@ -945,7 +958,7 @@ def start_background_job( # Outer nohup redirects to /dev/null since output goes to log files inside sh -c bg_cmd = f"nohup sh -c {quoted_sh_command} < /dev/null > /dev/null 2>&1 &" - self.execute_command(sandbox_id, bg_cmd, timeout=30) + self.execute_command(sandbox_id, bg_cmd, timeout=30, user=user) return BackgroundJob( job_id=job_id, @@ -1632,11 +1645,17 @@ async def execute_command( working_dir: Optional[str] = None, env: Optional[Dict[str, str]] = None, timeout: Optional[int] = None, + user: Optional[str] = None, ) -> CommandResponse: """Execute command directly via gateway (async).""" auth = await self._auth_cache.get_or_refresh(sandbox_id) if await self._auth_cache.is_vm(sandbox_id): + if user is not None: + raise ValueError( + "The 'user' parameter is only supported for container sandboxes, " + "not VM sandboxes." + ) return await self._execute_command_connect_rpc( sandbox_id=sandbox_id, command=command, @@ -1653,6 +1672,7 @@ async def execute_command( working_dir=working_dir, env=env, timeout=timeout, + user=user, ) async def _execute_command_connect_rpc( @@ -1733,6 +1753,7 @@ async def _execute_command_rest( working_dir: Optional[str] = None, env: Optional[Dict[str, str]] = None, timeout: Optional[int] = None, + user: Optional[str] = None, ) -> CommandResponse: gateway_url = auth["gateway_url"].rstrip("/") url = f"{gateway_url}/{auth['user_ns']}/{auth['job_id']}/exec" @@ -1746,6 +1767,8 @@ async def _execute_command_rest( "sandbox_id": sandbox_id, "timeout": effective_timeout, } + if user is not None: + payload["user"] = user for attempt in range(MAX_409_RETRIES): try: @@ -1811,6 +1834,7 @@ async def start_background_job( command: str, working_dir: Optional[str] = None, env: Optional[Dict[str, str]] = None, + user: Optional[str] = None, ) -> BackgroundJob: """Start a long-running command in the background (async). @@ -1822,6 +1846,8 @@ async def start_background_job( command: Command to execute working_dir: Working directory for command execution env: Environment variables + user: Run the job as this user, like ``docker exec -u`` (username or + numeric UID, optionally USER:GROUP). Container sandboxes only. Returns: BackgroundJob with job_id and file paths for polling @@ -1856,7 +1882,7 @@ async def start_background_job( # Outer nohup redirects to /dev/null since output goes to log files inside sh -c bg_cmd = f"nohup sh -c {quoted_sh_command} < /dev/null > /dev/null 2>&1 &" - await self.execute_command(sandbox_id, bg_cmd, timeout=30) + await self.execute_command(sandbox_id, bg_cmd, timeout=30, user=user) return BackgroundJob( job_id=job_id, diff --git a/packages/prime/src/prime_cli/commands/sandbox.py b/packages/prime/src/prime_cli/commands/sandbox.py index bf9442d8..f2761ef2 100644 --- a/packages/prime/src/prime_cli/commands/sandbox.py +++ b/packages/prime/src/prime_cli/commands/sandbox.py @@ -1088,6 +1088,13 @@ def run( "--timeout", help="Timeout for the command in seconds", ), + user: Optional[str] = typer.Option( + None, + "-u", + "--user", + help="Run the command as this user (username or UID, optionally USER:GROUP), " + "like 'docker exec -u'. Container sandboxes only.", + ), ) -> None: """Execute a command in a sandbox. @@ -1128,6 +1135,8 @@ def run( console.print(f"[bold blue]Environment:[/bold blue] {obfuscated_env}") if timeout is not None: console.print(f"[bold blue]Timeout:[/bold blue] {timeout}s") + if user: + console.print(f"[bold blue]User:[/bold blue] {user}") start_time = time.perf_counter() @@ -1138,6 +1147,7 @@ def run( working_dir, env_vars if env_vars else None, timeout=timeout, + user=user, ) # End timing