picod: initialize and validate workspace when --workspace is set

docs/getting-started.md

@@ -144,6 +144,8 @@ spec:
   maxSessionDuration: "8h"
 ```
 
+**Note:** When picod starts with the `--workspace` flag, it automatically creates the workspace directory if it does not exist. The process working directory is also changed to this workspace, ensuring that relative file paths are resolved correctly within the workspace.
+
 Verify the CodeInterpreter is created:
 
 ```bash

example/code-interpreter/code-interpreter.yaml

@@ -12,7 +12,8 @@ spec:
     image: ghcr.io/volcano-sh/picod:latest
     imagePullPolicy: IfNotPresent
     args:
-      - --workspace=/root
+      # The workspace directory will be created automatically by picod at startup
+      - --workspace=/workspace
     resources:
       limits:
         cpu: "500m"

integrations/dify-plugin/provider/agentcube.py

@@ -19,7 +19,7 @@
 
 
 class AgentcubeCodeInterpreterProvider(ToolProvider):
-    
+
     def _validate_credentials(self, credentials: dict[str, Any]) -> None:
         try:
             """
@@ -43,7 +43,7 @@ def _validate_credentials(self, credentials: dict[str, Any]) -> None:
     #     except Exception as e:
     #         raise ToolProviderOAuthError(str(e))
     #     return ""
-        
+
     # def _oauth_get_credentials(
     #     self, redirect_uri: str, system_credentials: Mapping[str, Any], request: Request
     # ) -> Mapping[str, Any]:

integrations/dify-plugin/tools/agentcube-code-interpreter.py

@@ -24,9 +24,13 @@ class AgentcubeCodeInterpreterTool(Tool):
     def _invoke(self, tool_parameters: dict[str, Any]) -> Generator[ToolInvokeMessage]:
         result = self.execute(**tool_parameters)
         yield self.create_json_message(result)
-
 
-    def execute(self, router_url=None, workload_manager_url=None, language="python", code_interpreter_id=None, session_id=None, code=None, command=None, session_reuse=False, **kwargs):
+
+    def execute(
+        self, router_url=None, workload_manager_url=None, language="python",
+        code_interpreter_id=None, session_id=None, code=None,
+        command=None, session_reuse=False, **kwargs,
+    ):
         # Validate required URLs
         if not router_url or not workload_manager_url:
             return {"status": "error", "reason": "router_url and workload_manager_url are required"}
@@ -49,11 +53,11 @@ def execute(self, router_url=None, workload_manager_url=None, language="python",
             if command:
                 command_result = ci_client.execute_command(command)
                 results.append({"type": "command", "result": command_result})
-            
+
             if language and code:
                 code_result = ci_client.run_code(language, code)
                 results.append({"type": "code", "result": code_result})
-            
+
             if not command and not code:
                 raise ValueError("Either command or code must be provided")
         except Exception as e:
@@ -71,11 +75,10 @@ def execute(self, router_url=None, workload_manager_url=None, language="python",
                 result["session_id"] = ci_client.session_id
         else:
             result = {
-                "status": "success", 
-                "session_id": ci_client.session_id, 
+                "status": "success",
+                "session_id": ci_client.session_id,
                 "code_interpreter_id": ci_client.name,
                 "results": results
             }
-        
+
         return result
-

pkg/picod/execute.go

@@ -102,6 +102,16 @@ func (s *Server) ExecuteHandler(c *gin.Context) {
 			})
 			return
 		}
+
+		// Ensure working directory exists
+		if err := os.MkdirAll(safeWorkingDir, 0755); err != nil {
+			c.JSON(http.StatusInternalServerError, gin.H{
+				"error": fmt.Sprintf("failed to create working directory: %v", err),
+				"code":  http.StatusInternalServerError,
+			})
+			return
+		}
+
 		cmd.Dir = safeWorkingDir
 	}
 

pkg/picod/files.go

@@ -389,14 +389,44 @@ func parseFileMode(modeStr string) os.FileMode {
 // setWorkspace sets the global workspace directory
 func (s *Server) setWorkspace(dir string) {
 	klog.Infof("setWorkspace called with dir: %q", dir)
+
+
+	// Save original working directory before changing it (only once)
+	if s.originalWorkingDir == "" {
+		cwd, err := os.Getwd()
+		if err != nil {
+			klog.Warningf("failed to get current working directory: %v", err)
+		} else {
+			s.originalWorkingDir = cwd
+		}
+	}
+
+	// Resolve to absolute path
+
 	absDir, err := filepath.Abs(dir)
 	if err != nil {
-		klog.Warningf("Failed to resolve absolute path for workspace '%s': %v", dir, err)
-		s.workspaceDir = dir // Fallback to provided path
-	} else {
-		s.workspaceDir = absDir
-		klog.Infof("Resolved workspace to absolute path: %q", s.workspaceDir)
+		klog.Fatalf("failed to resolve absolute path for workspace %q: %v", dir, err)
 	}
+
+	// Create directory if it doesn't exist
+	if err := os.MkdirAll(absDir, 0755); err != nil {
+		klog.Fatalf("failed to create workspace directory %q: %v", absDir, err)
+	}
+
+
+	// Verify path exists and is a directory
+	stat, err := os.Stat(absDir)
+	if err != nil {
+		klog.Fatalf("failed to stat workspace directory %q: %v", absDir, err)
+	}
+	if !stat.IsDir() {
+		klog.Fatalf("workspace path %q is not a directory", absDir)
+	}
+
+	// Set workspace directory
+	s.workspaceDir = absDir
+
+	klog.Infof("workspace directory initialized: %q", s.workspaceDir)
 }
 
 // sanitizePath ensures path is within allowed scope, preventing directory traversal attacks

pkg/picod/picod_test.go

@@ -82,6 +82,15 @@ func setupTestServer(t *testing.T, pubPEM string) (*Server, *httptest.Server, st
 }
 
 func TestPicoD_EndToEnd(t *testing.T) {
+	// Capture current working directory and restore it in cleanup
+	originalWd, err := os.Getwd()
+	require.NoError(t, err)
+	t.Cleanup(func() {
+		if err := os.Chdir(originalWd); err != nil {
+			t.Logf("failed to restore working directory: %v", err)
+		}
+	})
+
 	// 1. Setup Keys - single key pair for Router-style auth
 	routerPriv, routerPubStr := generateRSAKeys(t)
 
@@ -92,11 +101,8 @@ func TestPicoD_EndToEnd(t *testing.T) {
 	defer os.Unsetenv(PublicKeyEnvVar)
 
 	// Switch to temp dir for relative path tests
-	originalWd, err := os.Getwd()
-	require.NoError(t, err)
 	err = os.Chdir(tmpDir)
 	require.NoError(t, err)
-	defer func() { require.NoError(t, os.Chdir(originalWd)) }()
 
 	client := ts.Client()
 
@@ -172,25 +178,55 @@ func TestPicoD_EndToEnd(t *testing.T) {
 		assert.Equal(t, 124, resp.ExitCode)
 		assert.Contains(t, resp.Stderr, "Command timed out")
 
-		// 5. Working Directory Escape (Should Fail)
+		// 5. Non-existent Working Directory (Should Create It)
+		nonExistReq := ExecuteRequest{
+			Command:    []string{"sh", "-c", "pwd"},
+			WorkingDir: "subdir/nested",
+		}
+		nonExistBody, _ := json.Marshal(nonExistReq)
+		nonExistClaims := jwt.MapClaims{
+			"iat": time.Now().Unix(),
+			"exp": time.Now().Add(time.Hour * 6).Unix(),
+		}
+		nonExistToken := createToken(t, routerPriv, nonExistClaims)
+
+		nonExistReqHTTP, _ := http.NewRequest("POST", ts.URL+"/api/execute", bytes.NewBuffer(nonExistBody))
+		nonExistReqHTTP.Header.Set("Authorization", "Bearer "+nonExistToken)
+		nonExistReqHTTP.Header.Set("Content-Type", "application/json")
+
+		httpResp, err := client.Do(nonExistReqHTTP)
+		require.NoError(t, err)
+		assert.Equal(t, http.StatusOK, httpResp.StatusCode)
+
+		var nonExistResp ExecuteResponse
+		err = json.NewDecoder(httpResp.Body).Decode(&nonExistResp)
+		require.NoError(t, err)
+		assert.Equal(t, 0, nonExistResp.ExitCode)
+		assert.NotEmpty(t, nonExistResp.Stdout)
+
+		// Verify directory was created
+		_, err = os.Stat("subdir/nested")
+		assert.NoError(t, err)
+
+		// 6. Working Directory Escape (Should Fail)
 		escapeReq := ExecuteRequest{
 			Command:    []string{"ls"},
 			WorkingDir: "../",
 		}
 		escapeBody, _ := json.Marshal(escapeReq)
-		claims := jwt.MapClaims{
+		escapeClaims := jwt.MapClaims{
 			"iat": time.Now().Unix(),
 			"exp": time.Now().Add(time.Hour * 6).Unix(),
 		}
-		token := createToken(t, routerPriv, claims)
+		escapeToken := createToken(t, routerPriv, escapeClaims)
 
-		req, _ := http.NewRequest("POST", ts.URL+"/api/execute", bytes.NewBuffer(escapeBody))
-		req.Header.Set("Authorization", "Bearer "+token)
-		req.Header.Set("Content-Type", "application/json")
+		escapeReqHTTP, _ := http.NewRequest("POST", ts.URL+"/api/execute", bytes.NewBuffer(escapeBody))
+		escapeReqHTTP.Header.Set("Authorization", "Bearer "+escapeToken)
+		escapeReqHTTP.Header.Set("Content-Type", "application/json")
 
-		httpResp, err := client.Do(req)
+		escapeResp, err := client.Do(escapeReqHTTP)
 		require.NoError(t, err)
-		assert.Equal(t, http.StatusBadRequest, httpResp.StatusCode)
+		assert.Equal(t, http.StatusBadRequest, escapeResp.StatusCode)
 	})
 
 	t.Run("File Operations", func(t *testing.T) {
@@ -328,17 +364,23 @@ func TestPicoD_EndToEnd(t *testing.T) {
 // requires public key at startup. Without it, PicoD will fail to start.
 
 func TestPicoD_DefaultWorkspace(t *testing.T) {
+	// Capture current working directory and restore it in cleanup
+	originalWd, err := os.Getwd()
+	require.NoError(t, err)
+	t.Cleanup(func() {
+		if err := os.Chdir(originalWd); err != nil {
+			t.Logf("failed to restore working directory: %v", err)
+		}
+	})
+
 	// Setup temporary directory for test
 	tmpDir, err := os.MkdirTemp("", "picod_default_workspace_test")
 	require.NoError(t, err)
 	defer os.RemoveAll(tmpDir)
 
 	// Switch to temp dir
-	originalWd, err := os.Getwd()
-	require.NoError(t, err)
 	err = os.Chdir(tmpDir)
 	require.NoError(t, err)
-	defer func() { require.NoError(t, os.Chdir(originalWd)) }()
 
 	// Set public key env
 	_, pubStr := generateRSAKeys(t)
@@ -352,6 +394,7 @@ func TestPicoD_DefaultWorkspace(t *testing.T) {
 	}
 
 	server := NewServer(config)
+	defer server.RestoreWorkingDirectory()
 
 	// Verify workspaceDir is set to current working directory
 	cwd, err := os.Getwd()
@@ -364,6 +407,15 @@ func TestPicoD_DefaultWorkspace(t *testing.T) {
 }
 
 func TestPicoD_SetWorkspace(t *testing.T) {
+	// Capture current working directory and restore it in cleanup
+	originalWd, err := os.Getwd()
+	require.NoError(t, err)
+	t.Cleanup(func() {
+		if err := os.Chdir(originalWd); err != nil {
+			t.Logf("failed to restore working directory: %v", err)
+		}
+	})
+
 	// Setup temp dir
 	tmpDir, err := os.MkdirTemp("", "picod_setworkspace_test")
 	require.NoError(t, err)
@@ -401,11 +453,8 @@ func TestPicoD_SetWorkspace(t *testing.T) {
 	assert.Equal(t, resolve(absPath), resolve(server.workspaceDir))
 
 	// Case 2: Relative Path
-	originalWd, err := os.Getwd()
-	require.NoError(t, err)
 	err = os.Chdir(tmpDir)
 	require.NoError(t, err)
-	defer func() { require.NoError(t, os.Chdir(originalWd)) }()
 
 	server.setWorkspace("real")
 	assert.Equal(t, resolve(absPath), resolve(server.workspaceDir))

pkg/picod/server.go

@@ -34,11 +34,12 @@ type Config struct {
 
 // Server defines the PicoD HTTP server
 type Server struct {
-	engine       *gin.Engine
-	config       Config
-	authManager  *AuthManager
-	startTime    time.Time
-	workspaceDir string
+	engine             *gin.Engine
+	config             Config
+	authManager        *AuthManager
+	startTime          time.Time
+	workspaceDir       string
+	originalWorkingDir string
 }
 
 // NewServer creates a new PicoD server instance
@@ -110,6 +111,17 @@ func (s *Server) Run() error {
 	return server.ListenAndServe()
 }
 
+// RestoreWorkingDirectory restores the process working directory to its original state
+func (s *Server) RestoreWorkingDirectory() {
+	if s.originalWorkingDir != "" {
+		if err := os.Chdir(s.originalWorkingDir); err != nil {
+			klog.Warningf("failed to restore working directory to %q: %v", s.originalWorkingDir, err)
+		} else {
+			klog.Infof("restored working directory to: %q", s.originalWorkingDir)
+		}
+	}
+}
+
 // HealthCheckHandler handles health check requests
 func (s *Server) HealthCheckHandler(c *gin.Context) {
 	c.JSON(http.StatusOK, gin.H{