Skip to content

Commit aad049b

Browse files
dplakonoz-agent
dplakon
and
oz-agent
authored
feat: add --target-dir flag to direct backend (#37)
Co-authored-by: Oz <oz-agent@warp.dev>
1 parent 06b463e commit aad049b

4 files changed

Lines changed: 93 additions & 24 deletions

File tree

‎internal/config/config.go‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -41,6 +41,7 @@ type DockerConfig struct {
4141
// DirectConfig holds direct-backend-specific configuration.
4242
type DirectConfig struct {
4343
WorkspaceRoot string `yaml:"workspace_root"`
44+
TargetDir string `yaml:"target_dir"`
4445
OzPath string `yaml:"oz_path"`
4546
SetupCommand string `yaml:"setup_command"`
4647
TeardownCommand string `yaml:"teardown_command"`

‎internal/config/config_test.go‎

Lines changed: 25 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -235,6 +235,31 @@ backend:
235235
}
236236
}
237237

238+
func TestLoadDirectConfigWithTargetDir(t *testing.T) {
239+
path := writeTestConfig(t, `
240+
worker_id: "direct-worker"
241+
backend:
242+
direct:
243+
target_dir: "/home/user/myrepo"
244+
setup_command: "/opt/setup.sh"
245+
`)
246+
247+
cfg, err := Load(path)
248+
if err != nil {
249+
t.Fatalf("unexpected error: %v", err)
250+
}
251+
252+
if cfg.Backend.Direct == nil {
253+
t.Fatal("expected direct backend to be set")
254+
}
255+
if cfg.Backend.Direct.TargetDir != "/home/user/myrepo" {
256+
t.Errorf("target_dir = %q, want %q", cfg.Backend.Direct.TargetDir, "/home/user/myrepo")
257+
}
258+
if cfg.Backend.Direct.WorkspaceRoot != "" {
259+
t.Errorf("workspace_root should be empty, got %q", cfg.Backend.Direct.WorkspaceRoot)
260+
}
261+
}
262+
238263
func TestLoadBothBackendsError(t *testing.T) {
239264
path := writeTestConfig(t, `
240265
backend:

‎internal/worker/direct.go‎

Lines changed: 58 additions & 23 deletions
Original file line numberDiff line numberDiff line change
@@ -34,6 +34,7 @@ func hostBaseEnv() []string {
3434
// DirectBackendConfig holds configuration specific to the direct (non-containerized) backend.
3535
type DirectBackendConfig struct {
3636
WorkspaceRoot string
37+
TargetDir string // If set, run all tasks in this directory instead of creating per-task workspaces.
3738
OzPath string // Path to the oz CLI binary. If empty, looks up "oz" in PATH.
3839
SetupCommand string
3940
TeardownCommand string
@@ -59,13 +60,25 @@ func NewDirectBackend(ctx context.Context, config DirectBackendConfig) (*DirectB
5960
}
6061
log.Infof(ctx, "Using oz CLI at: %s", ozPath)
6162

62-
if config.WorkspaceRoot == "" {
63-
config.WorkspaceRoot = defaultWorkspaceRoot
64-
}
63+
if config.TargetDir != "" {
64+
// Validate that the target directory exists.
65+
info, err := os.Stat(config.TargetDir)
66+
if err != nil {
67+
return nil, fmt.Errorf("target directory %s does not exist: %w", config.TargetDir, err)
68+
}
69+
if !info.IsDir() {
70+
return nil, fmt.Errorf("target directory %s is not a directory", config.TargetDir)
71+
}
72+
log.Infof(ctx, "Using shared target directory: %s (per-task workspace isolation disabled)", config.TargetDir)
73+
} else {
74+
if config.WorkspaceRoot == "" {
75+
config.WorkspaceRoot = defaultWorkspaceRoot
76+
}
6577

66-
// Ensure workspace root exists.
67-
if err := os.MkdirAll(config.WorkspaceRoot, 0755); err != nil {
68-
return nil, fmt.Errorf("failed to create workspace root %s: %w", config.WorkspaceRoot, err)
78+
// Ensure workspace root exists.
79+
if err := os.MkdirAll(config.WorkspaceRoot, 0755); err != nil {
80+
return nil, fmt.Errorf("failed to create workspace root %s: %w", config.WorkspaceRoot, err)
81+
}
6982
}
7083

7184
return &DirectBackend{
@@ -78,14 +91,27 @@ func NewDirectBackend(ctx context.Context, config DirectBackendConfig) (*DirectB
7891
func (b *DirectBackend) ExecuteTask(ctx context.Context, params *TaskParams) error {
7992
taskID := params.TaskID
8093

81-
// 1. Create per-task workspace directory.
82-
workspaceDir := filepath.Join(b.config.WorkspaceRoot, taskID)
83-
if err := os.MkdirAll(workspaceDir, 0755); err != nil {
84-
return fmt.Errorf("failed to create workspace directory: %w", err)
94+
// Determine working directory: shared target dir or per-task workspace.
95+
var workspaceDir string
96+
usingTargetDir := b.config.TargetDir != ""
97+
98+
if usingTargetDir {
99+
workspaceDir = b.config.TargetDir
100+
} else {
101+
// Create per-task workspace directory.
102+
workspaceDir = filepath.Join(b.config.WorkspaceRoot, taskID)
103+
if err := os.MkdirAll(workspaceDir, 0755); err != nil {
104+
return fmt.Errorf("failed to create workspace directory: %w", err)
105+
}
106+
log.Infof(ctx, "Created workspace: %s", workspaceDir)
85107
}
86-
log.Infof(ctx, "Created workspace: %s", workspaceDir)
87108

88109
defer func() {
110+
if usingTargetDir {
111+
// Don't clean up the shared target directory.
112+
b.runTeardownIfConfigured(ctx, taskID, workspaceDir)
113+
return
114+
}
89115
if b.config.NoCleanup {
90116
log.Infof(ctx, "Skipping cleanup for workspace: %s", workspaceDir)
91117
return
@@ -102,6 +128,7 @@ func (b *DirectBackend) ExecuteTask(ctx context.Context, params *TaskParams) err
102128
if err := envFile.Close(); err != nil {
103129
return fmt.Errorf("failed to close environment file: %w", err)
104130
}
131+
defer os.Remove(envFilePath)
105132

106133
// 3. Build environment variables: common + config-level.
107134
envVars := make([]string, len(params.EnvVars))
@@ -159,6 +186,9 @@ func (b *DirectBackend) ExecuteTask(ctx context.Context, params *TaskParams) err
159186

160187
// Shutdown cleans up any workspace directories left behind under the workspace root.
161188
func (b *DirectBackend) Shutdown(ctx context.Context) {
189+
if b.config.WorkspaceRoot == "" {
190+
return
191+
}
162192
entries, err := os.ReadDir(b.config.WorkspaceRoot)
163193
if err != nil {
164194
if !os.IsNotExist(err) {
@@ -176,20 +206,25 @@ func (b *DirectBackend) Shutdown(ctx context.Context) {
176206
}
177207
}
178208

209+
// runTeardownIfConfigured runs the teardown command if one is configured.
210+
func (b *DirectBackend) runTeardownIfConfigured(ctx context.Context, taskID, workspaceDir string) {
211+
if b.config.TeardownCommand == "" {
212+
return
213+
}
214+
teardownEnv := []string{
215+
fmt.Sprintf("OZ_WORKSPACE_ROOT=%s", workspaceDir),
216+
"OZ_WORKER_BACKEND=direct",
217+
fmt.Sprintf("OZ_RUN_ID=%s", taskID),
218+
}
219+
log.Infof(ctx, "Running teardown command: %s", b.config.TeardownCommand)
220+
if err := b.runCommand(ctx, b.config.TeardownCommand, workspaceDir, teardownEnv); err != nil {
221+
log.Warnf(ctx, "Teardown command failed: %v", err)
222+
}
223+
}
224+
179225
// cleanup runs the teardown command (if configured) and removes the workspace directory.
180226
func (b *DirectBackend) cleanup(ctx context.Context, taskID, workspaceDir string) {
181-
if b.config.TeardownCommand != "" {
182-
teardownEnv := []string{
183-
fmt.Sprintf("OZ_WORKSPACE_ROOT=%s", workspaceDir),
184-
"OZ_WORKER_BACKEND=direct",
185-
fmt.Sprintf("OZ_RUN_ID=%s", taskID),
186-
}
187-
188-
log.Infof(ctx, "Running teardown command: %s", b.config.TeardownCommand)
189-
if err := b.runCommand(ctx, b.config.TeardownCommand, workspaceDir, teardownEnv); err != nil {
190-
log.Warnf(ctx, "Teardown command failed: %v", err)
191-
}
192-
}
227+
b.runTeardownIfConfigured(ctx, taskID, workspaceDir)
193228

194229
log.Infof(ctx, "Removing workspace: %s", workspaceDir)
195230
if err := os.RemoveAll(workspaceDir); err != nil {

‎main.go‎

Lines changed: 9 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -22,6 +22,7 @@ var CLI struct {
2222
WebSocketURL string `default:"wss://oz.warp.dev/api/v1/selfhosted/worker/ws" hidden:""`
2323
ServerRootURL string `default:"https://app.warp.dev" hidden:""`
2424
LogLevel string `help:"Log level (debug, info, warn, error)" default:"info" enum:"debug,info,warn,error"`
25+
TargetDir string `help:"Run all tasks in this directory instead of creating per-task workspaces (direct backend only)"`
2526
NoCleanup bool `help:"Do not remove containers after execution (for debugging)"`
2627
Volumes []string `help:"Volume mounts for task containers (format: HOST_PATH:CONTAINER_PATH or HOST_PATH:CONTAINER_PATH:MODE)" short:"v"`
2728
Env []string `help:"Environment variables for task containers (format: KEY=VALUE or KEY to pass through from host)" short:"e"`
@@ -152,16 +153,23 @@ func mergeConfig(fileConfig *config.FileConfig) (worker.Config, error) {
152153
mergedEnv[k] = v
153154
}
154155

155-
var workspaceRoot, ozPath, setupCmd, teardownCmd string
156+
var workspaceRoot, targetDir, ozPath, setupCmd, teardownCmd string
156157
if fileConfig != nil && fileConfig.Backend.Direct != nil {
157158
workspaceRoot = fileConfig.Backend.Direct.WorkspaceRoot
159+
targetDir = fileConfig.Backend.Direct.TargetDir
158160
ozPath = fileConfig.Backend.Direct.OzPath
159161
setupCmd = fileConfig.Backend.Direct.SetupCommand
160162
teardownCmd = fileConfig.Backend.Direct.TeardownCommand
161163
}
162164

165+
// CLI --target-dir overrides config file.
166+
if CLI.TargetDir != "" {
167+
targetDir = CLI.TargetDir
168+
}
169+
163170
wc.Direct = &worker.DirectBackendConfig{
164171
WorkspaceRoot: workspaceRoot,
172+
TargetDir: targetDir,
165173
OzPath: ozPath,
166174
SetupCommand: setupCmd,
167175
TeardownCommand: teardownCmd,

0 commit comments

Comments
 (0)