clone complete OMP agent state for OAuth

Author: lemon-mint

README.md

@@ -255,13 +255,12 @@ enable only the integrations that project needs.
 OMP is opt-in. When enabled, its selected binary, native modules, model catalog,
 and non-secret configuration are cloned into workspace-owned writable snapshots;
 container writes cannot modify the host copies. Set `requireCow: true` to require
-reflink support instead of allowing a private copy-once fallback. Set
-`agents.omp.import.oauthDB: true` only when the workspace may receive credentials.
-Cohotfs directly byte-copies `agent.db` and any existing `-wal`, `-shm`, or
-`-journal` sidecars into a private, writable snapshot; these files do not use
-reflinks even when `requireCow: true`. Host and workspace credential updates do
-not sync. Stop OMP or otherwise quiesce its database before workspace creation
-when a point-in-time-consistent copy is required.
+reflink support instead of allowing a private copy-once fallback. Setting
+`agents.omp.import.oauthDB: true` clones the complete OMP agent directory,
+including its credential, session, history, configuration, and cache state, into
+a private writable COW snapshot mounted at `PI_CODING_AGENT_DIR`. Host and
+workspace changes do not sync. Stop OMP before workspace creation when a
+point-in-time-consistent SQLite snapshot is required.
 
 All host-side state stays under:
 

integration/docker/end_to_end_test.go

@@ -139,18 +139,32 @@ func TestWorkspaceEndToEnd(t *testing.T) {
 		}}
 		record, _, source := harness.createWorkspaceWithOptions(t, "it-reserved-mask", "manual", true, config.ResourceSpec{}, options)
 		record = harness.start(t, record)
-		after, err := os.Lstat(filepath.Join(source, ".omp"))
+		assertReservedMask := func() {
+			after, err := os.Lstat(filepath.Join(source, ".omp"))
+			if err != nil {
+				t.Fatal(err)
+			}
+			if !os.SameFile(before, after) {
+				t.Fatal("Docker replaced the host reserved directory")
+			}
+			assertFile(t, filepath.Join(source, ".omp", "host-sentinel"), "host-only\n")
+			stdout, stderr := harness.ssh(t, record, nil, `test -d /workspace/.omp && test ! -e /workspace/.omp/host-sentinel && printf 'reserved-mask-ok\n'`)
+			if string(stdout) != "reserved-mask-ok\n" || len(stderr) != 0 {
+				t.Fatalf("reserved mask stdout=%q stderr=%q", stdout, stderr)
+			}
+		}
+		assertReservedMask()
+		record = harness.stop(t, record)
+		afterStop, err := os.Lstat(filepath.Join(source, ".omp"))
 		if err != nil {
 			t.Fatal(err)
 		}
-		if !os.SameFile(before, after) {
-			t.Fatal("Docker replaced the host reserved directory")
+		if !os.SameFile(before, afterStop) {
+			t.Fatal("stopping replaced the host reserved directory")
 		}
 		assertFile(t, filepath.Join(source, ".omp", "host-sentinel"), "host-only\n")
-		stdout, stderr := harness.ssh(t, record, nil, `test -d /workspace/.omp && test ! -e /workspace/.omp/host-sentinel && printf 'reserved-mask-ok\n'`)
-		if string(stdout) != "reserved-mask-ok\n" || len(stderr) != 0 {
-			t.Fatalf("reserved mask stdout=%q stderr=%q", stdout, stderr)
-		}
+		record = harness.start(t, record)
+		assertReservedMask()
 		harness.remove(t, record)
 	})
 	t.Run("go-toolchain-with-omp-oauth", func(t *testing.T) {

internal/ompimport/plan_linux.go

@@ -1,8 +1,7 @@
 //go:build linux
 
-// Package ompimport compiles explicit host OMP file imports. Selected host
-// files, including the OAuth database and its existing sidecars, are copied
-// into workspace-owned writable snapshots.
+// Package ompimport compiles explicit host OMP imports into workspace-owned
+// writable snapshots.
 package ompimport
 
 import (
@@ -113,16 +112,22 @@ func Compile(root *hostroot.Root, workspaceID string, spec config.OMPAgentSpec,
 		}
 		plan.Mounts = append(plan.Mounts, runtime.Mount{Source: snapshot, Target: containerNative, Type: "bind", Propagation: "rprivate"})
 	}
-	if spec.Import.OAuthDB && sources.Agent == "" {
-		return Plan{}, fmt.Errorf("OMP OAuth database source is unavailable")
-	}
-	if (spec.Import.Models || spec.Import.Config || spec.Import.OAuthDB) && sources.Agent != "" {
+	if spec.Import.OAuthDB {
+		if sources.Agent == "" {
+			return Plan{}, fmt.Errorf("OMP agent directory source is unavailable")
+		}
+		snapshot, err := prepareDirectorySnapshot(root, workspaceID, "agent", sources.Agent, spec.Import.RequireCOW)
+		if err != nil {
+			return Plan{}, err
+		}
+		plan.Mounts = append(plan.Mounts, runtime.Mount{Source: snapshot, Target: containerAgent, Type: "bind", Propagation: "rprivate"})
+	} else if (spec.Import.Models || spec.Import.Config) && sources.Agent != "" {
 		selected, err := selectedAgentFiles(spec.Import, sources.Agent)
 		if err != nil {
 			return Plan{}, err
 		}
-		if len(selected) != 0 || spec.Import.OAuthDB {
-			snapshot, err := prepareAgentSnapshot(root, workspaceID, sources.Agent, selected, spec.Import.RequireCOW)
+		if len(selected) != 0 {
+			snapshot, err := prepareSelectedSnapshot(root, workspaceID, "agent", sources.Agent, selected, spec.Import.RequireCOW)
 			if err != nil {
 				return Plan{}, err
 			}
@@ -141,19 +146,6 @@ func selectedAgentFiles(spec config.OMPImportSpec, root string) ([]string, error
 		candidates = append(candidates, "models.yml", "models.yaml")
 	}
 	selected := make([]string, 0, len(candidates))
-	if spec.OAuthDB {
-		for _, name := range []string{"agent.db", "agent.db-wal", "agent.db-shm", "agent.db-journal"} {
-			path := filepath.Join(root, name)
-			err := requireRegular(path, false)
-			if name != "agent.db" && errors.Is(err, os.ErrNotExist) {
-				continue
-			}
-			if err != nil {
-				return nil, fmt.Errorf("validate OMP OAuth database file %s: %w", name, err)
-			}
-			selected = append(selected, name)
-		}
-	}
 	for _, name := range candidates {
 		path := filepath.Join(root, name)
 		err := requireRegular(path, false)
@@ -196,38 +188,6 @@ func prepareSelectedSnapshot(root *hostroot.Root, workspaceID, name, source stri
 	})
 }
 
-func prepareAgentSnapshot(root *hostroot.Root, workspaceID, source string, selected []string, requireCOW bool) (string, error) {
-	fingerprint, err := selectedFingerprint(source, selected)
-	if err != nil {
-		return "", err
-	}
-	return prepareSnapshot(root, workspaceID, "agent", fingerprint, func(staging string) error {
-		for _, relative := range selected {
-			sourcePath := filepath.Join(source, relative)
-			destinationPath := filepath.Join(staging, relative)
-			var err error
-			if isOAuthDatabaseFile(relative) {
-				err = copyRegularFile(sourcePath, destinationPath, false)
-			} else {
-				err = cloneRegularFile(sourcePath, destinationPath, false, requireCOW)
-			}
-			if err != nil {
-				return err
-			}
-		}
-		return nil
-	})
-}
-
-func isOAuthDatabaseFile(name string) bool {
-	switch name {
-	case "agent.db", "agent.db-wal", "agent.db-shm", "agent.db-journal":
-		return true
-	default:
-		return false
-	}
-}
-
 func prepareDirectorySnapshot(root *hostroot.Root, workspaceID, name, source string, requireCOW bool) (string, error) {
 	fingerprint, err := directoryFingerprint(source)
 	if err != nil {
@@ -318,10 +278,6 @@ func cloneRegularFile(source, destination string, executable, requireCOW bool) e
 	return writeRegularFile(source, destination, executable, true, requireCOW)
 }
 
-func copyRegularFile(source, destination string, executable bool) error {
-	return writeRegularFile(source, destination, executable, false, false)
-}
-
 func writeRegularFile(source, destination string, executable, tryReflink, requireCOW bool) error {
 	sourceFD, err := unix.Open(source, unix.O_RDONLY|unix.O_NOFOLLOW|unix.O_CLOEXEC, 0)
 	if err != nil {
@@ -399,13 +355,16 @@ func selectedFingerprint(root string, names []string) (string, error) {
 	sort.Strings(names)
 	for _, name := range names {
 		path := filepath.Join(root, name)
-		entry, err := os.Stat(path)
+		entry, err := os.Lstat(path)
 		if err != nil {
 			return "", err
 		}
 		if _, err := fmt.Fprintf(hash, "%s\x00%d\x00%d\x00%d\x00", name, entry.Mode(), entry.Size(), entry.ModTime().UnixNano()); err != nil {
 			return "", err
 		}
+		if err := fingerprintRegularFile(hash, path); err != nil {
+			return "", err
+		}
 	}
 	return hex.EncodeToString(hash.Sum(nil)[:16]), nil
 }
@@ -426,9 +385,30 @@ func fingerprintEntry(hash interface{ Write([]byte) (int, error) }, path, relati
 		_, err = fmt.Fprintf(hash, "%s\x00", target)
 		return err
 	}
+	if info.Mode().IsRegular() {
+		return fingerprintRegularFile(hash, path)
+	}
 	return nil
 }
 
+func fingerprintRegularFile(hash io.Writer, path string) error {
+	fd, err := unix.Open(path, unix.O_RDONLY|unix.O_NOFOLLOW|unix.O_CLOEXEC, 0)
+	if err != nil {
+		return err
+	}
+	file := os.NewFile(uintptr(fd), path)
+	defer file.Close()
+	var stat unix.Stat_t
+	if err := unix.Fstat(fd, &stat); err != nil {
+		return err
+	}
+	if stat.Mode&unix.S_IFMT != unix.S_IFREG {
+		return fmt.Errorf("unsafe OMP file %s", path)
+	}
+	_, err = io.Copy(hash, file)
+	return err
+}
+
 func requireRegular(path string, executable bool) error {
 	info, err := os.Lstat(path)
 	if err != nil {

internal/ompimport/plan_linux_test.go

@@ -68,14 +68,21 @@ func TestCompileBuildsWritableReflinkMountsForSelectedOMPState(t *testing.T) {
 	}
 }
 
-func TestCompileCopiesOAuthDatabaseFilesDespiteRequireCOW(t *testing.T) {
+func TestCompileClonesCompleteAgentDirectoryWhenOAuthEnabled(t *testing.T) {
 	root := openTestRoot(t)
 	defer root.Close()
 	sources := ompFixture(t)
+	nested := filepath.Join(sources.Agent, "sessions", "current")
+	if err := os.MkdirAll(nested, 0o700); err != nil {
+		t.Fatal(err)
+	}
+	if err := os.WriteFile(filepath.Join(nested, "session.json"), []byte("session fixture"), 0o600); err != nil {
+		t.Fatal(err)
+	}
 	spec := config.OMPAgentSpec{
 		Enabled: true,
 		Import: config.OMPImportSpec{
-			Enabled: true, OAuthDB: true, RequireCOW: true,
+			Enabled: true, OAuthDB: true, RequireCOW: false,
 		},
 	}
 
@@ -84,7 +91,10 @@ func TestCompileCopiesOAuthDatabaseFilesDespiteRequireCOW(t *testing.T) {
 		t.Fatal(err)
 	}
 	agent := requiredMount(t, plan, containerAgent).Source
-	for _, name := range []string{"agent.db", "agent.db-wal", "agent.db-shm", "agent.db-journal"} {
+	for _, name := range []string{
+		"agent.db", "agent.db-wal", "agent.db-shm", "agent.db-journal",
+		"config.yml", "models.yml", "history.db", filepath.Join("sessions", "current", "session.json"),
+	} {
 		sourcePath := filepath.Join(sources.Agent, name)
 		snapshotPath := filepath.Join(agent, name)
 		source, err := os.ReadFile(sourcePath)
@@ -96,7 +106,7 @@ func TestCompileCopiesOAuthDatabaseFilesDespiteRequireCOW(t *testing.T) {
 			t.Fatal(err)
 		}
 		if string(snapshot) != string(source) {
-			t.Fatalf("OAuth database file %s changed during copy", name)
+			t.Fatalf("OMP agent file %s changed during clone", name)
 		}
 		sourceInfo, err := os.Stat(sourcePath)
 		if err != nil {
@@ -107,7 +117,7 @@ func TestCompileCopiesOAuthDatabaseFilesDespiteRequireCOW(t *testing.T) {
 			t.Fatal(err)
 		}
 		if os.SameFile(sourceInfo, snapshotInfo) {
-			t.Fatalf("OAuth database file %s reuses the host inode", name)
+			t.Fatalf("OMP agent file %s reuses the host inode", name)
 		}
 	}
 	if err := os.WriteFile(filepath.Join(agent, "agent.db"), []byte("workspace-update"), 0o600); err != nil {
@@ -208,6 +218,59 @@ func TestCompileChangesSnapshotWhenHostStateChanges(t *testing.T) {
 	}
 }
 
+func TestCompileInvalidatesFullAgentSnapshotWhenContentChangesWithoutMetadata(t *testing.T) {
+	root := openTestRoot(t)
+	defer root.Close()
+	sources := ompFixture(t)
+	spec := config.OMPAgentSpec{
+		Enabled: true,
+		Import:  config.OMPImportSpec{Enabled: true, OAuthDB: true},
+	}
+	first, err := Compile(root, testWorkspaceID, spec, sources)
+	if err != nil {
+		t.Fatal(err)
+	}
+	historyPath := filepath.Join(sources.Agent, "history.db")
+	info, err := os.Stat(historyPath)
+	if err != nil {
+		t.Fatal(err)
+	}
+	changed, err := os.ReadFile(historyPath)
+	if err != nil {
+		t.Fatal(err)
+	}
+	changed[0] ^= 0xff
+	if err := os.WriteFile(historyPath, changed, info.Mode().Perm()); err != nil {
+		t.Fatal(err)
+	}
+	if err := os.Chtimes(historyPath, info.ModTime(), info.ModTime()); err != nil {
+		t.Fatal(err)
+	}
+	changedInfo, err := os.Stat(historyPath)
+	if err != nil {
+		t.Fatal(err)
+	}
+	if changedInfo.Size() != info.Size() || !changedInfo.ModTime().Equal(info.ModTime()) {
+		t.Fatalf("test mutation changed metadata: before=%#v after=%#v", info, changedInfo)
+	}
+	second, err := Compile(root, testWorkspaceID, spec, sources)
+	if err != nil {
+		t.Fatal(err)
+	}
+	firstAgent := requiredMount(t, first, containerAgent).Source
+	secondAgent := requiredMount(t, second, containerAgent).Source
+	if firstAgent == secondAgent {
+		t.Fatal("content change with stable size and mtime reused the previous full agent snapshot")
+	}
+	snapshot, err := os.ReadFile(filepath.Join(secondAgent, "history.db"))
+	if err != nil {
+		t.Fatal(err)
+	}
+	if string(snapshot) != string(changed) {
+		t.Fatal("new full agent snapshot does not contain changed content")
+	}
+}
+
 func TestDiscoverUsesExplicitAgentDirectoryAndPrivateNatives(t *testing.T) {
 	home := t.TempDir()
 	bin := filepath.Join(t.TempDir(), "bin")

internal/workspace/docker.go

@@ -705,7 +705,15 @@ func (s *DockerService) startLocked(ctx context.Context, record state.Workspace)
 	}
 	releasePins, err := pinReservedWorkspaceMasks(plan)
 	if err != nil {
-		return record, err
+		if !status.Running {
+			return record, err
+		}
+		cleanupCtx := context.WithoutCancel(ctx)
+		stopErr := s.backend.Stop(cleanupCtx, record.RuntimeRef, 10*time.Second)
+		leaseErr := s.releaseActiveIntegrationLeases(cleanupCtx, &record)
+		transitionErr := record.Transition(state.StatusError, s.now())
+		saveErr := s.store.SaveWorkspace(record)
+		return record, errors.Join(err, stopErr, leaseErr, transitionErr, saveErr)
 	}
 	defer func() { releasePins() }()
 	if err := s.acquireIntegrationLeases(ctx, &record, plan); err != nil {