diff --git a/docs/getting-started.md b/docs/getting-started.md index 485e7a18..92da77af 100644 --- a/docs/getting-started.md +++ b/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 diff --git a/example/code-interpreter/code-interpreter.yaml b/example/code-interpreter/code-interpreter.yaml index 96020cbe..c63649cb 100644 --- a/example/code-interpreter/code-interpreter.yaml +++ b/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" diff --git a/integrations/dify-plugin/provider/agentcube.py b/integrations/dify-plugin/provider/agentcube.py index 82eed36d..430d8cee 100644 --- a/integrations/dify-plugin/provider/agentcube.py +++ b/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]: diff --git a/integrations/dify-plugin/tools/agentcube-code-interpreter.py b/integrations/dify-plugin/tools/agentcube-code-interpreter.py index 9bc5a3fd..f912aae3 100644 --- a/integrations/dify-plugin/tools/agentcube-code-interpreter.py +++ b/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 - diff --git a/pkg/picod/execute.go b/pkg/picod/execute.go index b6c4dd10..ee644f07 100644 --- a/pkg/picod/execute.go +++ b/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 } diff --git a/pkg/picod/files.go b/pkg/picod/files.go index 24830009..c0cb081d 100644 --- a/pkg/picod/files.go +++ b/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 diff --git a/pkg/picod/picod_test.go b/pkg/picod/picod_test.go index efe6f6eb..6f3d1c79 100644 --- a/pkg/picod/picod_test.go +++ b/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)) diff --git a/pkg/picod/server.go b/pkg/picod/server.go index 0da8f8b2..1218e1dd 100644 --- a/pkg/picod/server.go +++ b/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{