feat: implement Supervisor Orchestration & Mocking (#11)

* feat: implement MountWorkspace and CopyOut in Supervisor

* docs: update roadmap for orchestration and mocking completion

* refactor: optimize locking in MountWorkspace

* feat: add policy validation to CopyOut

* test: enhance supervisor test robustness and state coverage

Author: yanurag-dev

ROADMAP.md

@@ -17,12 +17,12 @@ This roadmap tracks the progress of the Sandforge Agent Sandbox based on [ARCHIT
 ## Phase 2: Orchestration & Mocking
 *Goal: Build the state machine that manages sandbox lifecycles.*
 
-- [ ] **2.1 Sandbox Supervisor**:
-    - [ ] Implementation of the Lifecycle State Machine (Requested -> Provisioning -> Ready -> ...).
-    - [ ] Concurrent session management.
-- [ ] **2.2 Mock Backend Driver**:
-    - [ ] An in-memory/process-based driver for testing the supervisor without a VM.
-- [ ] **2.3 Artifact Manager**: Basic logic to handle "CopyOut" for logs and files.
+- [x] **2.1 Sandbox Supervisor**:
+    - [x] Implementation of the Lifecycle State Machine (Start, Run, Stop, Mount, CopyOut).
+    - [x] Concurrent session management.
+- [x] **2.2 Mock Backend Driver**:
+    - [x] An in-memory/process-based driver for testing the supervisor without a VM.
+- [x] **2.3 Artifact Manager**: Basic logic to handle "CopyOut" for logs and files.
 
 ## Phase 3: macOS Execution Plane (macos-vz)
 *Goal: Boot a real Linux VM on macOS using the Apple Virtualization Framework.*

internal/policy/engine.go

@@ -27,18 +27,30 @@ type Engine struct {
 }
 
 func (e *Engine) EvaluateMount(mount api.WorkspaceMount) error {
-	path := filepath.Clean(mount.HostPath)
+	return e.EvaluateHostPath(mount.HostPath)
+}
+
+func (e *Engine) EvaluateHostPath(path string) error {
+	path = filepath.Clean(path)
 
 	if !filepath.IsAbs(path) {
 		return ErrPathNotAbs
 	}
 
 	// Resolve symlinks to prevent bypasses (e.g., a symlink pointing to /etc)
+	// For CopyOut, the file might not exist yet, so we resolve the parent directory if the path itself doesn't exist.
 	resolved, err := filepath.EvalSymlinks(path)
 	if err != nil {
-		return err // Path must exist to be validated
+		// If path doesn't exist, try resolving the directory
+		dir := filepath.Dir(path)
+		resolvedDir, errDir := filepath.EvalSymlinks(dir)
+		if errDir != nil {
+			return errDir
+		}
+		path = filepath.Join(resolvedDir, filepath.Base(path))
+	} else {
+		path = resolved
 	}
-	path = resolved
 
 	allowed := false
 	for _, prefix := range e.AllowedHostPrefixes {
@@ -68,6 +80,7 @@ func (e *Engine) EvaluateMount(mount api.WorkspaceMount) error {
 	return nil
 }
 
+
 func (e *Engine) EvaluateSandbox(spec api.SandboxSpec) error {
 	if spec.CPU > e.MaxCPU {
 		return ErrResourceLimitExceeded

internal/supervisor/supervisor.go

@@ -13,14 +13,14 @@ import (
 type State string
 
 const (
-	StateRequested          State = "requested"
-	StateProvisioning       State = "provisioning"
-	StateReady              State = "ready"
-	StateExecuting          State = "executing"
-	StateCopyingArtifacts   State = "copying_artifacts"
-	StateDestroying         State = "destroying"
-	StateDestroyed          State = "destroyed"
-	StateError              State = "error"
+	StateRequested        State = "requested"
+	StateProvisioning     State = "provisioning"
+	StateReady            State = "ready"
+	StateExecuting        State = "executing"
+	StateCopyingArtifacts State = "copying_artifacts"
+	StateDestroying       State = "destroying"
+	StateDestroyed        State = "destroyed"
+	StateError            State = "error"
 )
 
 // SandboxInstance tracks the runtime state of a single sandbox.
@@ -210,3 +210,72 @@ func (s *Supervisor) Stop(id string) error {
 
 	return nil
 }
+
+// MountWorkspace allows attaching a host path to the sandbox.
+func (s *Supervisor) MountWorkspace(id string, mount api.WorkspaceMount) error {
+	// 1. Find the instance
+	s.mu.RLock()
+	instance, exists := s.instances[id]
+	s.mu.RUnlock()
+
+	if !exists {
+		return errors.New("sandbox not found")
+	}
+
+	// 2. Initial state check
+	instance.mu.RLock()
+	if instance.State != StateReady {
+		instance.mu.RUnlock()
+		return fmt.Errorf("sandbox is in state %s, must be %s to mount", instance.State, StateReady)
+	}
+	instance.mu.RUnlock()
+
+	// 3. Evaluate policy (without holding instance lock)
+	if err := s.policy.EvaluateMount(mount); err != nil {
+		return err
+	}
+
+	// 4. Re-check state and get handle
+	instance.mu.Lock()
+	if instance.State != StateReady {
+		instance.mu.Unlock()
+		return fmt.Errorf("sandbox is in state %s, must be %s to mount", instance.State, StateReady)
+	}
+	handle := instance.Handle
+	instance.mu.Unlock()
+
+	// 5. Call backend
+	return s.backend.MountWorkspace(handle, mount)
+}
+
+// CopyOut retrieves a file or directory from the sandbox.
+func (s *Supervisor) CopyOut(id string, path string, dest string) error {
+	// 1. Find the instance
+	s.mu.RLock()
+	instance, exists := s.instances[id]
+	s.mu.RUnlock()
+
+	if !exists {
+		return errors.New("sandbox not found")
+	}
+
+	// 2. Policy check for destination host path
+	if err := s.policy.EvaluateHostPath(dest); err != nil {
+		return fmt.Errorf("policy denied copy out destination: %w", err)
+	}
+
+	// 3. Validate state
+	instance.mu.RLock()
+	handle := instance.Handle
+	state := instance.State
+	instance.mu.RUnlock()
+
+	// We allow CopyOut if it's Ready or Executing (e.g. streaming logs)
+	if state != StateReady && state != StateExecuting {
+		return fmt.Errorf("sandbox is in state %s, must be %s or %s to copy out", state, StateReady, StateExecuting)
+	}
+
+	// 4. Call backend
+	return s.backend.CopyOut(handle, path, dest)
+}
+

internal/supervisor/supervisor_test.go

@@ -2,6 +2,7 @@ package supervisor
 
 import (
 	"fmt"
+	"path/filepath"
 	"sync"
 	"testing"
 
@@ -126,3 +127,101 @@ func TestSupervisorLifecycle(t *testing.T) {
 		}
 	})
 }
+
+func TestSupervisorMountAndCopy(t *testing.T) {
+	mockBackend := backend.NewMockBackend()
+	
+	// Create a temp dir for allowed mounts
+	tmpDir := t.TempDir()
+	
+	engine := &policy.Engine{
+		AllowedHostPrefixes: []string{tmpDir},
+		MaxCPU:              4,
+		MaxMemoryMb:         4096,
+		MaxDiskGb:           10,
+		AllowedNetworkModes: []string{"offline"},
+	}
+	sup, err := NewSupervisor(mockBackend, engine)
+	if err != nil {
+		t.Fatalf("Failed to create supervisor: %v", err)
+	}
+
+	id := "test-mount"
+	spec := api.SandboxSpec{CPU: 1, MemoryMb: 512, DiskGb: 1, NetworkMode: "offline"}
+	
+	if err := sup.Start(id, spec); err != nil {
+		t.Fatalf("Failed to start sandbox: %v", err)
+	}
+	defer sup.Stop(id)
+
+	t.Run("ValidMount", func(t *testing.T) {
+		err := sup.MountWorkspace(id, api.WorkspaceMount{
+			HostPath:  tmpDir,
+			GuestPath: "/workspace",
+		})
+		if err != nil {
+			t.Errorf("Expected valid mount to succeed, got %v", err)
+		}
+	})
+
+	t.Run("InvalidMount_Path", func(t *testing.T) {
+		err := sup.MountWorkspace(id, api.WorkspaceMount{
+			HostPath:  "/etc",
+			GuestPath: "/workspace",
+		})
+		if err == nil {
+			t.Error("Expected mount to /etc to be blocked by policy")
+		}
+	})
+
+	t.Run("InvalidMount_State", func(t *testing.T) {
+		// Manually set state to Executing to test rejection
+		sup.mu.Lock()
+		instance := sup.instances[id]
+		sup.mu.Unlock()
+		
+		originalState := instance.GetState()
+		instance.SetState(StateExecuting)
+		defer instance.SetState(originalState)
+
+		err := sup.MountWorkspace(id, api.WorkspaceMount{
+			HostPath:  tmpDir,
+			GuestPath: "/workspace",
+		})
+		if err == nil {
+			t.Error("Expected mount to fail when state is Executing")
+		}
+	})
+
+	t.Run("CopyOut_Valid", func(t *testing.T) {
+		dest := filepath.Join(tmpDir, "log.txt")
+		err := sup.CopyOut(id, "/workspace/log.txt", dest)
+		if err != nil {
+			t.Errorf("Expected CopyOut to succeed, got %v", err)
+		}
+	})
+
+	t.Run("CopyOut_InvalidPath", func(t *testing.T) {
+		err := sup.CopyOut(id, "/workspace/log.txt", "/etc/shadow")
+		if err == nil {
+			t.Error("Expected CopyOut to /etc/shadow to be blocked by policy")
+		}
+	})
+
+	t.Run("CopyOut_InvalidState", func(t *testing.T) {
+		sup.mu.Lock()
+		instance := sup.instances[id]
+		sup.mu.Unlock()
+
+		originalState := instance.GetState()
+		instance.SetState(StateError)
+		defer instance.SetState(originalState)
+
+		err := sup.CopyOut(id, "/workspace/log.txt", filepath.Join(tmpDir, "error.txt"))
+		if err == nil {
+			t.Error("Expected CopyOut to fail when state is Error")
+		}
+	})
+}
+
+