fix: guard spawn initiation and restore metadata

This commit is contained in:
harshitsinghbhandari 2026-05-29 15:43:58 +05:30
parent 1fee0164e2
commit fb7dbbbb42
4 changed files with 47 additions and 1 deletions

View File

@ -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()

View File

@ -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

View File

@ -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)

View File

@ -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 {