mcp

Author: ianhodge

internal/worker/direct.go

@@ -0,0 +1,225 @@
+package worker
+
+import (
+	"bufio"
+	"context"
+	"fmt"
+	"os"
+	"os/exec"
+	"path/filepath"
+	"strings"
+
+	"github.com/warpdotdev/oz-agent-worker/internal/log"
+)
+
+const defaultWorkspaceRoot = "/var/lib/oz/workspaces"
+
+// DirectBackendConfig holds configuration specific to the direct (non-containerized) backend.
+type DirectBackendConfig struct {
+	WorkspaceRoot   string
+	OzPath          string // Path to the oz CLI binary. If empty, looks up "oz" in PATH.
+	SetupCommand    string
+	TeardownCommand string
+	NoCleanup       bool
+	Env             map[string]string
+}
+
+// DirectBackend executes tasks directly on the host without Docker.
+type DirectBackend struct {
+	config DirectBackendConfig
+	ozPath string // resolved path to the oz CLI
+}
+
+// NewDirectBackend creates a new direct backend, verifying the oz CLI is available.
+func NewDirectBackend(ctx context.Context, config DirectBackendConfig) (*DirectBackend, error) {
+	ozPath := config.OzPath
+	if ozPath == "" {
+		var err error
+		ozPath, err = exec.LookPath("oz")
+		if err != nil {
+			return nil, fmt.Errorf("oz CLI not found in PATH: %w", err)
+		}
+	}
+	log.Infof(ctx, "Using oz CLI at: %s", ozPath)
+
+	if config.WorkspaceRoot == "" {
+		config.WorkspaceRoot = defaultWorkspaceRoot
+	}
+
+	// Ensure workspace root exists.
+	if err := os.MkdirAll(config.WorkspaceRoot, 0755); err != nil {
+		return nil, fmt.Errorf("failed to create workspace root %s: %w", config.WorkspaceRoot, err)
+	}
+
+	return &DirectBackend{
+		config: config,
+		ozPath: ozPath,
+	}, nil
+}
+
+// ExecuteTask runs the agent directly on the host.
+func (b *DirectBackend) ExecuteTask(ctx context.Context, params *TaskParams) error {
+	taskID := params.TaskID
+
+	// 1. Create per-task workspace directory.
+	workspaceDir := filepath.Join(b.config.WorkspaceRoot, taskID)
+	if err := os.MkdirAll(workspaceDir, 0755); err != nil {
+		return fmt.Errorf("failed to create workspace directory: %w", err)
+	}
+	log.Infof(ctx, "Created workspace: %s", workspaceDir)
+
+	defer func() {
+		if b.config.NoCleanup {
+			log.Infof(ctx, "Skipping cleanup for workspace: %s", workspaceDir)
+			return
+		}
+		b.cleanup(ctx, taskID, workspaceDir)
+	}()
+
+	// 2. Create temp environment file for setup script to write to.
+	envFile, err := os.CreateTemp(workspaceDir, "oz-env-*")
+	if err != nil {
+		return fmt.Errorf("failed to create environment file: %w", err)
+	}
+	envFilePath := envFile.Name()
+	if err := envFile.Close(); err != nil {
+		return fmt.Errorf("failed to close environment file: %w", err)
+	}
+
+	// 3. Build environment variables: common + config-level.
+	envVars := make([]string, len(params.EnvVars))
+	copy(envVars, params.EnvVars)
+	for key, value := range b.config.Env {
+		envVars = append(envVars, fmt.Sprintf("%s=%s", key, value))
+	}
+
+	// 4. Run setup command if configured.
+	if b.config.SetupCommand != "" {
+		setupEnv := append(envVars,
+			fmt.Sprintf("OZ_WORKSPACE_ROOT=%s", workspaceDir),
+			"OZ_WORKER_BACKEND=direct",
+			fmt.Sprintf("OZ_RUN_ID=%s", taskID),
+			fmt.Sprintf("OZ_ENVIRONMENT_FILE=%s", envFilePath),
+		)
+
+		log.Infof(ctx, "Running setup command: %s", b.config.SetupCommand)
+		if err := b.runCommand(ctx, b.config.SetupCommand, workspaceDir, setupEnv); err != nil {
+			return fmt.Errorf("setup command failed: %w", err)
+		}
+	}
+
+	// 5. Parse environment file for KEY=VALUE pairs written by setup script.
+	setupScriptEnv, err := parseEnvFile(envFilePath)
+	if err != nil {
+		log.Warnf(ctx, "Failed to parse environment file: %v", err)
+	}
+	for key, value := range setupScriptEnv {
+		envVars = append(envVars, fmt.Sprintf("%s=%s", key, value))
+	}
+
+	// 6. Invoke oz CLI with base args.
+	// Task env vars override host env vars (e.g. WARP_API_KEY from the assignment
+	// must take precedence over the worker's own WARP_API_KEY).
+	cmd := exec.CommandContext(ctx, b.ozPath, params.BaseArgs...)
+	cmd.Dir = workspaceDir
+	cmd.Env = mergeEnvVars(os.Environ(), envVars)
+	cmd.Stdout = os.Stdout
+	cmd.Stderr = os.Stderr
+
+	log.Infof(ctx, "Running oz agent in workspace %s", workspaceDir)
+	log.Debugf(ctx, "Command: %s %s", b.ozPath, strings.Join(params.BaseArgs, " "))
+
+	if err := cmd.Run(); err != nil {
+		return fmt.Errorf("oz agent exited with error: %w", err)
+	}
+
+	log.Infof(ctx, "Task %s execution completed successfully", taskID)
+	return nil
+}
+
+// Shutdown is a no-op for the direct backend.
+func (b *DirectBackend) Shutdown(ctx context.Context) {}
+
+// cleanup runs the teardown command (if configured) and removes the workspace directory.
+func (b *DirectBackend) cleanup(ctx context.Context, taskID, workspaceDir string) {
+	if b.config.TeardownCommand != "" {
+		teardownEnv := []string{
+			fmt.Sprintf("OZ_WORKSPACE_ROOT=%s", workspaceDir),
+			"OZ_WORKER_BACKEND=direct",
+			fmt.Sprintf("OZ_RUN_ID=%s", taskID),
+		}
+
+		log.Infof(ctx, "Running teardown command: %s", b.config.TeardownCommand)
+		if err := b.runCommand(ctx, b.config.TeardownCommand, workspaceDir, teardownEnv); err != nil {
+			log.Warnf(ctx, "Teardown command failed: %v", err)
+		}
+	}
+
+	log.Infof(ctx, "Removing workspace: %s", workspaceDir)
+	if err := os.RemoveAll(workspaceDir); err != nil {
+		log.Warnf(ctx, "Failed to remove workspace %s: %v", workspaceDir, err)
+	}
+}
+
+// runCommand executes a shell command with the given working directory and environment.
+func (b *DirectBackend) runCommand(ctx context.Context, command, dir string, env []string) error {
+	cmd := exec.CommandContext(ctx, "/bin/sh", "-c", command)
+	cmd.Dir = dir
+	cmd.Env = mergeEnvVars(os.Environ(), env)
+	cmd.Stdout = os.Stdout
+	cmd.Stderr = os.Stderr
+	return cmd.Run()
+}
+
+// mergeEnvVars merges base and override env var slices (KEY=VALUE format).
+// Override entries take precedence over base entries with the same key.
+func mergeEnvVars(base, override []string) []string {
+	envMap := make(map[string]string, len(base)+len(override))
+	var keys []string
+
+	for _, entry := range base {
+		key, _, _ := strings.Cut(entry, "=")
+		if _, exists := envMap[key]; !exists {
+			keys = append(keys, key)
+		}
+		envMap[key] = entry
+	}
+
+	for _, entry := range override {
+		key, _, _ := strings.Cut(entry, "=")
+		if _, exists := envMap[key]; !exists {
+			keys = append(keys, key)
+		}
+		envMap[key] = entry
+	}
+
+	result := make([]string, 0, len(keys))
+	for _, key := range keys {
+		result = append(result, envMap[key])
+	}
+	return result
+}
+
+// parseEnvFile reads a file of KEY=VALUE lines, skipping empty lines and comments.
+func parseEnvFile(path string) (map[string]string, error) {
+	f, err := os.Open(path)
+	if err != nil {
+		return nil, err
+	}
+	defer f.Close()
+
+	result := make(map[string]string)
+	scanner := bufio.NewScanner(f)
+	for scanner.Scan() {
+		line := strings.TrimSpace(scanner.Text())
+		if line == "" || strings.HasPrefix(line, "#") {
+			continue
+		}
+		key, value, ok := strings.Cut(line, "=")
+		if !ok {
+			continue
+		}
+		result[strings.TrimSpace(key)] = strings.TrimSpace(value)
+	}
+	return result, scanner.Err()
+}