Skip to content

Commit ac76bf2

Browse files
cwlbraacwlbraa
authored
real container resume redux redux (#129)
this PR implements "real" container resume where we store each Run WithExec in the container LLB definition and shores up the base of that container by saving off InitialSourceDir in state. InitialSourceDir is a made consistent and re-hydratable by using dag.Host().Directory(forkRepoPath).AsGit().Ref(initialWorktreeHead).Tree(). because we can rely on that SHA always pointing at the exact right base content and the forkRepo always having that SHA, we can source the base of sourcedir without keeping a copy of the first commit on disk. 6/25: this PR does have a big piece of "magic" in it. we template in the .git file to be gitdir: repos/name/worktrees/id so that, on export, we don't constantly override reset the git history. Other stuff buried in here: env.apply no longer takes arguments that give the false impression it's gonna make a commit. it just takes newState (the container) env.run calls apply run_background still doesn't. to make those restore we gotta do some additional thinking around how we data model service persistence. additionally on 6/25 cc @aluzzardi: better error handling in the event of failed commands. we now no longer fail the withexec for failed commands, but still propagate error information to the git notes and the tool responses. this is necessary for better prompt engineering around endpoints. the agent was getting confused and using localhost to try to talk to things it had run_backgrounded. closes #115. * working Signed-off-by: Connor Braa <connor@dagger.io> * cleanup Signed-off-by: Connor Braa <connor@dagger.io> * clean up apply Signed-off-by: Connor Braa <connor@dagger.io> * shore up error handling in run command Signed-off-by: Connor Braa <connor@dagger.io> * assorted fixes: - stop saving initialSourceDir - its saved in the llb. - have environment_update rebuild the base from workdir - improve prompt engineering around internal/external ports - improve run_command error handling so that failed commands can be saved into the container state Signed-off-by: Connor Braa <connor@dagger.io> * resolve build failure post-merge Signed-off-by: Connor Braa <connor@dagger.io> * make tests pass Signed-off-by: Connor Braa <connor@dagger.io> * shoutout guillaume's tests, i totally forgot about this part of the magic Signed-off-by: Connor Braa <connor@dagger.io> * remove unnecessary git notes initialization Signed-off-by: Connor Braa <connor@dagger.io> --------- Signed-off-by: Connor Braa <connor@dagger.io>
1 parent 2b175c5 commit ac76bf2

File tree

6 files changed

+85
-41
lines changed

6 files changed

+85
-41
lines changed

‎environment/environment.go‎

Lines changed: 49 additions & 26 deletions
Original file line numberDiff line numberDiff line change
@@ -20,7 +20,8 @@ type EnvironmentInfo struct {
2020
Config *EnvironmentConfig
2121
State *State
2222

23-
ID string
23+
ID string
24+
//TODO(braa): I think we only need this for Export, now. remove and pass explicitly.
2425
Worktree string
2526
}
2627

@@ -35,7 +36,7 @@ type Environment struct {
3536
mu sync.RWMutex
3637
}
3738

38-
func New(ctx context.Context, dag *dagger.Client, id, title, worktree string) (*Environment, error) {
39+
func New(ctx context.Context, dag *dagger.Client, id, title, worktree string, initialSourceDir *dagger.Directory) (*Environment, error) {
3940
env := &Environment{
4041
EnvironmentInfo: &EnvironmentInfo{
4142
ID: id,
@@ -56,14 +57,14 @@ func New(ctx context.Context, dag *dagger.Client, id, title, worktree string) (*
5657
}
5758
}
5859

59-
container, err := env.buildBase(ctx)
60+
container, err := env.buildBase(ctx, initialSourceDir)
6061
if err != nil {
6162
return nil, err
6263
}
6364

6465
slog.Info("Creating environment", "id", env.ID, "workdir", env.Config.Workdir)
6566

66-
if err := env.apply(ctx, "Create environment", "Create the environment", "", container); err != nil {
67+
if err := env.apply(ctx, container); err != nil {
6768
return nil, err
6869
}
6970

@@ -132,7 +133,8 @@ func LoadInfo(ctx context.Context, id string, state []byte, worktree string) (*E
132133
return envInfo, nil
133134
}
134135

135-
func (env *Environment) apply(ctx context.Context, name, explanation, output string, newState *dagger.Container) error {
136+
func (env *Environment) apply(ctx context.Context, newState *dagger.Container) error {
137+
// TODO(braa): is this sync redundant with newState.ID?
136138
if _, err := newState.Sync(ctx); err != nil {
137139
return err
138140
}
@@ -173,11 +175,7 @@ func containerWithEnvAndSecrets(dag *dagger.Client, container *dagger.Container,
173175
return container, nil
174176
}
175177

176-
func (env *Environment) buildBase(ctx context.Context) (*dagger.Container, error) {
177-
sourceDir := env.dag.Host().Directory(env.Worktree, dagger.HostDirectoryOpts{
178-
NoCache: true,
179-
})
180-
178+
func (env *Environment) buildBase(ctx context.Context, baseSourceDir *dagger.Directory) (*dagger.Container, error) {
181179
container := env.dag.
182180
Container().
183181
From(env.Config.BaseImage).
@@ -215,7 +213,7 @@ func (env *Environment) buildBase(ctx context.Context) (*dagger.Container, error
215213
container = container.WithServiceBinding(service.Config.Name, service.svc)
216214
}
217215

218-
container = container.WithDirectory(".", sourceDir)
216+
container = container.WithDirectory(".", baseSourceDir)
219217

220218
return container, nil
221219
}
@@ -227,42 +225,67 @@ func (env *Environment) UpdateConfig(ctx context.Context, explanation string, ne
227225

228226
env.Config = newConfig
229227

230-
// Re-build the base image from the worktree
231-
container, err := env.buildBase(ctx)
228+
// Get current working directory from container to preserve changes
229+
currentWorkdir := env.container().Directory(env.Config.Workdir)
230+
231+
// Re-build the base image with the new config
232+
container, err := env.buildBase(ctx, currentWorkdir)
232233
if err != nil {
233234
return err
234235
}
235236

236-
if err := env.apply(ctx, "Update environment", explanation, "", container); err != nil {
237+
if err := env.apply(ctx, container); err != nil {
237238
return err
238239
}
239240

240241
return nil
241242
}
242243

243-
func (env *Environment) Run(ctx context.Context, explanation, command, shell string, useEntrypoint bool) (string, error) {
244+
func (env *Environment) Run(ctx context.Context, command, shell string, useEntrypoint bool) (string, error) {
244245
args := []string{}
245246
if command != "" {
246247
args = []string{shell, "-c", command}
247248
}
248249
newState := env.container().WithExec(args, dagger.ContainerWithExecOpts{
249250
UseEntrypoint: useEntrypoint,
251+
Expect: dagger.ReturnTypeAny, // Don't treat non-zero exit as error
250252
})
253+
254+
exitCode, err := newState.ExitCode(ctx)
255+
if err != nil {
256+
return "", fmt.Errorf("failed to get exit code: %w", err)
257+
}
258+
251259
stdout, err := newState.Stdout(ctx)
252260
if err != nil {
253-
var exitErr *dagger.ExecError
254-
if errors.As(err, &exitErr) {
255-
env.Notes.Add("$ %s\nexit %d\nstdout: %s\nstderr: %s\n\n", command, exitErr.ExitCode, exitErr.Stdout, exitErr.Stderr)
256-
return fmt.Sprintf("command failed with exit code %d.\nstdout: %s\nstderr: %s", exitErr.ExitCode, exitErr.Stdout, exitErr.Stderr), nil
257-
}
258-
return "", err
261+
return "", fmt.Errorf("failed to get stdout: %w", err)
259262
}
260-
env.Notes.Add("$ %s\n%s\n\n", command, stdout)
261263

262-
return stdout, nil
264+
stderr, err := newState.Stderr(ctx)
265+
if err != nil {
266+
return "", fmt.Errorf("failed to get stderr: %w", err)
267+
}
268+
269+
// Log the command execution with all details
270+
env.Notes.Add("$ %s\nexit %d\nstdout: %s\nstderr: %s\n\n", command, exitCode, stdout, stderr)
271+
272+
// Always apply the container state (preserving changes even on non-zero exit)
273+
if err := env.apply(ctx, newState); err != nil {
274+
return stdout, fmt.Errorf("failed to apply container state: %w", err)
275+
}
276+
277+
// Return combined output (stdout + stderr if there was stderr)
278+
combinedOutput := stdout
279+
if stderr != "" {
280+
if stdout != "" {
281+
combinedOutput += "\n"
282+
}
283+
combinedOutput += "stderr: " + stderr
284+
}
285+
return combinedOutput, nil
263286
}
264287

265-
func (env *Environment) RunBackground(ctx context.Context, explanation, command, shell string, ports []int, useEntrypoint bool) (EndpointMappings, error) {
288+
func (env *Environment) RunBackground(ctx context.Context, command, shell string, ports []int, useEntrypoint bool) (EndpointMappings, error) {
266289
args := []string{}
267290
if command != "" {
268291
args = []string{shell, "-c", command}
@@ -316,15 +339,15 @@ func (env *Environment) RunBackground(ctx context.Context, explanation, command,
316339
if err != nil {
317340
return nil, err
318341
}
319-
endpoint.External = externalEndpoint
342+
endpoint.HostExternal = externalEndpoint
320343

321344
internalEndpoint, err := svc.Endpoint(ctx, dagger.ServiceEndpointOpts{
322345
Port: port,
323346
})
324347
if err != nil {
325348
return nil, err
326349
}
327-
endpoint.Internal = internalEndpoint
350+
endpoint.EnvironmentInternal = internalEndpoint
328351
}
329352

330353
return endpoints, nil

‎environment/filesystem.go‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -32,7 +32,7 @@ func (s *Environment) FileRead(ctx context.Context, targetFile string, shouldRea
3232
}
3333

3434
func (s *Environment) FileWrite(ctx context.Context, explanation, targetFile, contents string) error {
35-
err := s.apply(ctx, "Write "+targetFile, explanation, "", s.container().WithNewFile(targetFile, contents))
35+
err := s.apply(ctx, s.container().WithNewFile(targetFile, contents))
3636
if err != nil {
3737
return fmt.Errorf("failed applying file write, skipping git propogation: %w", err)
3838
}
@@ -43,7 +43,7 @@ func (s *Environment) FileWrite(ctx context.Context, explanation, targetFile, co
4343
}
4444

4545
func (s *Environment) FileDelete(ctx context.Context, explanation, targetFile string) error {
46-
err := s.apply(ctx, "Delete "+targetFile, explanation, "", s.container().WithoutFile(targetFile))
46+
err := s.apply(ctx, s.container().WithoutFile(targetFile))
4747
if err != nil {
4848
return err
4949
}

‎environment/integration/helpers.go‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -217,7 +217,7 @@ func (u *UserActions) RunCommand(envID, command, explanation string) string {
217217
env, err := u.repo.Get(u.ctx, u.dag, envID)
218218
require.NoError(u.t, err, "Failed to get environment %s", envID)
219219

220-
output, err := env.Run(u.ctx, explanation, command, "/bin/sh", false)
220+
output, err := env.Run(u.ctx, command, "/bin/sh", false)
221221
require.NoError(u.t, err, "Run command should succeed")
222222

223223
err = u.repo.Update(u.ctx, env, "Run "+command, explanation)

‎environment/service.go‎

Lines changed: 5 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -21,8 +21,8 @@ type Service struct {
2121
}
2222

2323
type EndpointMapping struct {
24-
Internal string `json:"internal"`
25-
External string `json:"external"`
24+
EnvironmentInternal string `json:"environment_internal"`
25+
HostExternal string `json:"host_external"`
2626
}
2727

2828
type EndpointMappings map[int]*EndpointMapping
@@ -81,7 +81,7 @@ func (env *Environment) startService(ctx context.Context, cfg *ServiceConfig) (*
8181
endpoints := EndpointMappings{}
8282
for _, port := range cfg.ExposedPorts {
8383
endpoint := &EndpointMapping{
84-
Internal: fmt.Sprintf("%s:%d", cfg.Name, port),
84+
EnvironmentInternal: fmt.Sprintf("%s:%d", cfg.Name, port),
8585
}
8686
endpoints[port] = endpoint
8787

@@ -103,7 +103,7 @@ func (env *Environment) startService(ctx context.Context, cfg *ServiceConfig) (*
103103
if err != nil {
104104
return nil, fmt.Errorf("failed to get endpoint for service %s: %w", cfg.Name, err)
105105
}
106-
endpoint.External = externalEndpoint
106+
endpoint.HostExternal = externalEndpoint
107107
}
108108

109109
return &Service{
@@ -125,7 +125,7 @@ func (env *Environment) AddService(ctx context.Context, explanation string, cfg
125125
env.Services = append(env.Services, svc)
126126

127127
state := env.container().WithServiceBinding(cfg.Name, svc.svc)
128-
if err := env.apply(ctx, "Add service "+cfg.Name, explanation, "", state); err != nil {
128+
if err := env.apply(ctx, state); err != nil {
129129
return nil, err
130130
}
131131

‎mcpserver/tools.go‎

Lines changed: 6 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -412,7 +412,7 @@ var EnvironmentListTool = &Tool{
412412

413413
var EnvironmentRunCmdTool = &Tool{
414414
Definition: mcp.NewTool("environment_run_cmd",
415-
mcp.WithDescription("Run a command on behalf of the user in the terminal."),
415+
mcp.WithDescription("Run a terminal command inside the environment."),
416416
mcp.WithString("explanation",
417417
mcp.Description("One sentence explanation for why this command is being run."),
418418
),
@@ -440,7 +440,7 @@ Failure to do so will result in the tool being stuck, awaiting for the command t
440440
mcp.Description("Use the image entrypoint, if present, by prepending it to the args."),
441441
),
442442
mcp.WithArray("ports",
443-
mcp.Description("Ports to expose. Only works with background environments. For each port, returns the internal (for use by other environments) and external (for use by the user) address."),
443+
mcp.Description("Ports to expose. Only works with background environments. For each port, returns the environment_internal (for use inside environments) and host_external (for use by the user) addresses."),
444444
mcp.Items(map[string]any{"type": "number"}),
445445
),
446446
),
@@ -468,7 +468,7 @@ Failure to do so will result in the tool being stuck, awaiting for the command t
468468
ports = append(ports, int(port.(float64)))
469469
}
470470
}
471-
endpoints, runErr := env.RunBackground(ctx, request.GetString("explanation", ""), command, shell, ports, request.GetBool("use_entrypoint", false))
471+
endpoints, runErr := env.RunBackground(ctx, command, shell, ports, request.GetBool("use_entrypoint", false))
472472
// We want to update the repository even if the command failed.
473473
if resp, err := updateRepo(); err != nil {
474474
return resp, nil
@@ -490,13 +490,13 @@ Background commands are unaffected by filesystem and any other kind of changes.
490490
string(out), env.Config.Workdir, env.ID)), nil
491491
}
492492

493-
stdout, runErr := env.Run(ctx, request.GetString("explanation", ""), command, shell, request.GetBool("use_entrypoint", false))
493+
stdout, runErr := env.Run(ctx, command, shell, request.GetBool("use_entrypoint", false))
494494
// We want to update the repository even if the command failed.
495495
if resp, err := updateRepo(); err != nil {
496496
return resp, nil
497497
}
498498
if runErr != nil {
499-
return mcp.NewToolResultErrorFromErr("failed to run command", err), nil
499+
return mcp.NewToolResultErrorFromErr("failed to run command", runErr), nil
500500
}
501501

502502
return mcp.NewToolResultText(fmt.Sprintf("%s\n\nAny changes to the container workdir (%s) have been committed and pushed to container-use/ remote", stdout, env.Config.Workdir)), nil
@@ -744,7 +744,7 @@ var EnvironmentAddServiceTool = &Tool{
744744
mcp.Description("The command to start the service. If not provided the image default command will be used."),
745745
),
746746
mcp.WithArray("ports",
747-
mcp.Description("Ports to expose. For each port, returns the internal (for use by other environments) and external (for use by the user) address."),
747+
mcp.Description("Ports to expose. For each port, returns the environment_internal (for use by environments) and external (for use by the user) address."),
748748
mcp.Items(map[string]any{"type": "number"}),
749749
),
750750
mcp.WithArray("envs",

‎repository/repository.go‎

Lines changed: 22 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -157,7 +157,28 @@ func (r *Repository) Create(ctx context.Context, dag *dagger.Client, description
157157
return nil, err
158158
}
159159

160-
env, err := environment.New(ctx, dag, id, description, worktree)
160+
worktreeHead, err := RunGitCommand(ctx, worktree, "rev-parse", "HEAD")
161+
if err != nil {
162+
return nil, err
163+
}
164+
worktreeHead = strings.TrimSpace(worktreeHead)
165+
166+
baseSourceDir, err := dag.
167+
Host().
168+
Directory(r.forkRepoPath, dagger.HostDirectoryOpts{NoCache: true}). // bust cache for each Create call
169+
AsGit().
170+
Ref(worktreeHead).
171+
Tree(dagger.GitRefTreeOpts{DiscardGitDir: true}).
172+
WithNewFile(
173+
".git", // make .git point back at the fork repo TODO(braa): ideally we do this on export, but we'd need to move some of that logic into repository
174+
fmt.Sprintf("gitdir: %s/worktrees/%s", r.forkRepoPath, id),
175+
).
176+
Sync(ctx) // don't bust cache when loading from state
177+
if err != nil {
178+
return nil, fmt.Errorf("failed loading initial source directory: %w", err)
179+
}
180+
181+
env, err := environment.New(ctx, dag, id, description, worktree, baseSourceDir)
161182
if err != nil {
162183
return nil, err
163184
}

0 commit comments

Comments
 (0)