From fb7dbbbb420e088f02064f855c98786a4a2bf90d Mon Sep 17 00:00:00 2001
From: harshitsinghbhandari <24b4506@iitb.ac.in>
Date: Fri, 29 May 2026 15:43:58 +0530
Subject: [PATCH] fix: guard spawn initiation and restore metadata

---
 backend/internal/lifecycle/manager.go      |  6 +++++-
 backend/internal/lifecycle/manager_test.go | 19 +++++++++++++++++++
 backend/internal/session/manager.go        |  6 ++++++
 backend/internal/session/manager_test.go   | 17 +++++++++++++++++
 4 files changed, 47 insertions(+), 1 deletion(-)

diff --git a/backend/internal/lifecycle/manager.go b/backend/internal/lifecycle/manager.go
index c8e74f2ce..7494e346f 100644
--- a/backend/internal/lifecycle/manager.go
+++ b/backend/internal/lifecycle/manager.go
@@ -292,7 +292,11 @@ func (m *Manager) OnSpawnInitiated(ctx context.Context, rec domain.SessionRecord
 		if current, ok, err := m.store.Get(ctx, rec.ID); err != nil {
 			return err
 		} else if ok {
-			cur = current.Lifecycle
+			currentLC := current.Lifecycle
+			if !isTerminal(currentLC.Session.State) && !isTerminal(rec.Lifecycle.Session.State) {
+				return fmt.Errorf("lifecycle: OnSpawnInitiated for active session %q", rec.ID)
+			}
+			cur = currentLC
 		}
 		rec.Lifecycle = m.prepareLifecycleWrite(cur, rec.Lifecycle)
 		now := m.clock()
diff --git a/backend/internal/lifecycle/manager_test.go b/backend/internal/lifecycle/manager_test.go
index c5f75ad0f..fb7bf84d5 100644
--- a/backend/internal/lifecycle/manager_test.go
+++ b/backend/internal/lifecycle/manager_test.go
@@ -392,6 +392,25 @@ func TestOnSpawnCompleted(t *testing.T) {
 	}
 }
 
+func TestOnSpawnInitiated_ActiveSessionRejected(t *testing.T) {
+	mgr, store := newManager()
+	store.seed(sid, lc(domain.SessionWorking, domain.ReasonTaskInProgress, domain.RuntimeAlive))
+
+	err := mgr.OnSpawnInitiated(context.Background(), domain.SessionRecord{
+		ID:        sid,
+		ProjectID: domain.ProjectID("proj"),
+		Lifecycle: lc(domain.SessionNotStarted, domain.ReasonSpawnRequested, domain.RuntimeUnknown),
+	})
+	if err == nil {
+		t.Fatal("OnSpawnInitiated should reject a non-terminal row on top of an active session")
+	}
+
+	got := mustLoad(t, store)
+	if got.Session.State != domain.SessionWorking || got.Revision != 0 {
+		t.Fatalf("active row should be unchanged, got %+v", got)
+	}
+}
+
 func TestOnKillRequested(t *testing.T) {
 	tests := []struct {
 		name        string
diff --git a/backend/internal/session/manager.go b/backend/internal/session/manager.go
index 6ae2a60e9..c7ccf3396 100644
--- a/backend/internal/session/manager.go
+++ b/backend/internal/session/manager.go
@@ -320,9 +320,15 @@ func (m *Manager) Restore(ctx context.Context, id domain.SessionID) (domain.Sess
 	}
 	if err := m.lcm.OnSpawnCompleted(ctx, id, outcome); err != nil {
 		m.rollbackRuntime(ctx, handle)
+		// Re-upsert the original record to undo the reopen.
 		if revertErr := m.lcm.OnSpawnInitiated(ctx, rec); revertErr != nil {
 			return domain.Session{}, fmt.Errorf("restore %s: revert after spawn completed failure: %w (original error: %v)", id, revertErr, err)
 		}
+		if len(rec.Metadata) > 0 {
+			if revertErr := m.store.PatchMetadata(ctx, id, rec.Metadata); revertErr != nil {
+				return domain.Session{}, fmt.Errorf("restore %s: revert metadata after spawn completed failure: %w (original error: %v)", id, revertErr, err)
+			}
+		}
 		return domain.Session{}, fmt.Errorf("restore %s: on spawn completed: %w", id, err)
 	}
 	return m.Get(ctx, id)
diff --git a/backend/internal/session/manager_test.go b/backend/internal/session/manager_test.go
index a0812456d..b54cf4b09 100644
--- a/backend/internal/session/manager_test.go
+++ b/backend/internal/session/manager_test.go
@@ -473,6 +473,7 @@ func TestRestore_OnSpawnCompletedFailure_RollsBackRuntime(t *testing.T) {
 	if err := h.store.PatchMetadata(ctx, "sess-1", map[string]string{lifecycle.MetaAgentSessionID: "agent-xyz"}); err != nil {
 		t.Fatalf("patch metadata: %v", err)
 	}
+	beforeMeta, _ := h.store.GetMetadata(ctx, "sess-1")
 
 	// Fail the post-create LCM call; capture teardown counts just before restore.
 	h.lcm.onSpawnErr = errors.New("lcm boom")
@@ -491,6 +492,10 @@ func TestRestore_OnSpawnCompletedFailure_RollsBackRuntime(t *testing.T) {
 	if rec.Lifecycle.Revision != before.Lifecycle.Revision+2 {
 		t.Fatalf("restore failure should advance revision twice, got %d want %d", rec.Lifecycle.Revision, before.Lifecycle.Revision+2)
 	}
+	afterMeta, _ := h.store.GetMetadata(ctx, "sess-1")
+	if !equalStringMap(afterMeta, beforeMeta) {
+		t.Fatalf("restore failure should restore metadata, got %+v want %+v", afterMeta, beforeMeta)
+	}
 
 	// The runtime created during restore is torn back down so no process is
 	// stranded; the workspace is left intact (it holds the agent's prior work).
@@ -572,6 +577,18 @@ func equalStrings(a, b []string) bool {
 	return true
 }
 
+func equalStringMap(a, b map[string]string) bool {
+	if len(a) != len(b) {
+		return false
+	}
+	for k, v := range a {
+		if b[k] != v {
+			return false
+		}
+	}
+	return true
+}
+
 func contains(ids []domain.SessionID, id domain.SessionID) bool {
 	for _, x := range ids {
 		if x == id {