feat(orchestrator): persist ACP session id in session metadata (#1748)

* feat(orchestrator): persist ACP session id in session metadata storeResumeToken now mirrors the agent's ACP session id into the session's durable acp metadata map after the CAS succeeds. Previously the id only reached metadata when an agent emitted a session_info frame, so most sessions could not be joined to their agent-CLI transcripts once their executors_running row was pruned. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> * fix(orchestrator): make ACP session id mirror atomic and rotation-safe Address review feedback on the metadata mirror: replace the read-modify-write in persistACPSessionID with a single guarded UPDATE (SetSessionACPSessionID). json_patch merges the session_id sub-key atomically, preserving session_info keys, and the write is gated on executors_running still holding the id as its resume token — so a delayed event from a rotated execution can never overwrite the live execution's id, and the session_info TOCTOU disappears along with the unlogged GetTaskSession read. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> --------- Co-authored-by: Kandev Agent <agent@kandev.dev> Co-authored-by: Claude Fable 5 <noreply@anthropic.com>
2026-07-17 01:54:04 +01:00
parent 1248121fd5
commit 934b713ee9
6 changed files with 165 additions and 0 deletions
@@ -69,6 +69,12 @@ func (s *Service) storeResumeToken(ctx context.Context, taskID, sessionID, expec
 			zap.String("expected_exec_id", expectedExecID),
 			zap.String("resume_token", acpSessionID),
 			zap.String("last_message_uuid", lastMessageUUID))
+		// The CAS above proved the token belongs to the session's current
+		// execution — mirror it into the session's durable "acp" metadata
+		// too. executors_running rows are operational state and get pruned;
+		// without this, the ACP session id only survives for agents that
+		// happen to emit a session_info frame (handleSessionInfoEvent).
+		s.persistACPSessionID(ctx, sessionID, acpSessionID)
 	case errors.Is(err, models.ErrExecutionRotated):
 		// The row's agent_execution_id has rotated since the event was emitted —
 		// the token belongs to a defunct execution and must be dropped. The new
@@ -96,3 +102,24 @@ func (s *Service) storeResumeToken(ctx context.Context, taskID, sessionID, expec
 			zap.Error(err))
 	}
 }
+
+// persistACPSessionID mirrors the agent's ACP session id into the session's
+// "acp" metadata map. Best-effort: resume correctness never depends on this
+// copy — it exists so the id survives executors_running cleanup for consumers
+// that join sessions to agent-CLI artifacts (e.g. transcript-based usage
+// stats). The repo performs the mirror as a single guarded UPDATE: it merges
+// into the existing acp map (preserving session_info keys), skips when the
+// stored id is already current, and only writes while executors_running still
+// holds acpSessionID as the resume token — so a stale event from a rotated
+// execution can never overwrite the live execution's id.
+func (s *Service) persistACPSessionID(ctx context.Context, sessionID, acpSessionID string) {
+	if sessionID == "" || acpSessionID == "" || s.repo == nil {
+		return
+	}
+	if _, err := s.repo.SetSessionACPSessionID(ctx, sessionID, acpSessionID); err != nil {
+		s.logger.Warn("failed to persist ACP session id to session metadata",
+			zap.String("session_id", sessionID),
+			zap.String("acp_session_id", acpSessionID),
+			zap.Error(err))
+	}
+}
@@ -0,0 +1,101 @@
+package orchestrator
+
+import (
+	"context"
+	"testing"
+
+	"github.com/stretchr/testify/require"
+
+	"github.com/kandev/kandev/internal/task/models"
+	sqliterepo "github.com/kandev/kandev/internal/task/repository/sqlite"
+)
+
+// seedResumeToken puts the session's executors_running row into the state
+// storeResumeToken leaves it in after a successful CAS: resume_token holds the
+// agent's ACP session id. persistACPSessionID only mirrors ids that match this
+// row, so every subtest that expects a write needs it.
+func seedResumeToken(t *testing.T, repo *sqliterepo.Repository, sessionID, taskID, token string) {
+	t.Helper()
+	require.NoError(t, repo.UpsertExecutorRunning(context.Background(), &models.ExecutorRunning{
+		ID:          sessionID,
+		SessionID:   sessionID,
+		TaskID:      taskID,
+		ResumeToken: token,
+		Status:      "ready",
+	}))
+}
+
+// persistACPSessionID must write the ACP session id into the session's "acp"
+// metadata map, preserve keys session_info already stored, be a no-op when the
+// stored id is current, and refuse to mirror an id that no longer matches the
+// executors_running resume token (stale event from a rotated execution).
+func TestPersistACPSessionID(t *testing.T) {
+	ctx := context.Background()
+
+	t.Run("writes id into empty metadata", func(t *testing.T) {
+		repo := setupTestRepo(t)
+		seedSession(t, repo, "t1", "s1", "step1")
+		seedResumeToken(t, repo, "s1", "t1", "acp-123")
+		svc := createTestService(repo, newMockStepGetter(), newMockTaskRepo())
+
+		svc.persistACPSessionID(ctx, "s1", "acp-123")
+
+		session, err := repo.GetTaskSession(ctx, "s1")
+		require.NoError(t, err)
+		acp, ok := session.Metadata["acp"].(map[string]interface{})
+		require.True(t, ok, "acp metadata map missing: %+v", session.Metadata)
+		require.Equal(t, "acp-123", acp["session_id"])
+	})
+
+	t.Run("merges with existing acp map and updates stale id", func(t *testing.T) {
+		repo := setupTestRepo(t)
+		seedSession(t, repo, "t1", "s1", "step1")
+		seedResumeToken(t, repo, "s1", "t1", "acp-new")
+		require.NoError(t, repo.SetSessionMetadataKey(ctx, "s1", "acp",
+			map[string]interface{}{"session_id": "acp-old", "title": "My session"}))
+		svc := createTestService(repo, newMockStepGetter(), newMockTaskRepo())
+
+		svc.persistACPSessionID(ctx, "s1", "acp-new")
+
+		session, err := repo.GetTaskSession(ctx, "s1")
+		require.NoError(t, err)
+		acp := session.Metadata["acp"].(map[string]interface{})
+		require.Equal(t, "acp-new", acp["session_id"])
+		require.Equal(t, "My session", acp["title"], "session_info keys must survive")
+	})
+
+	t.Run("skips write when id is current", func(t *testing.T) {
+		repo := setupTestRepo(t)
+		seedSession(t, repo, "t1", "s1", "step1")
+		seedResumeToken(t, repo, "s1", "t1", "acp-123")
+		require.NoError(t, repo.SetSessionMetadataKey(ctx, "s1", "acp",
+			map[string]interface{}{"session_id": "acp-123"}))
+		before, err := repo.GetTaskSession(ctx, "s1")
+		require.NoError(t, err)
+
+		svc := createTestService(repo, newMockStepGetter(), newMockTaskRepo())
+		svc.persistACPSessionID(ctx, "s1", "acp-123")
+
+		after, err := repo.GetTaskSession(ctx, "s1")
+		require.NoError(t, err)
+		require.Equal(t, before.UpdatedAt, after.UpdatedAt, "no-op must not touch the row")
+	})
+
+	t.Run("skips write when execution rotated to a newer token", func(t *testing.T) {
+		repo := setupTestRepo(t)
+		seedSession(t, repo, "t1", "s1", "step1")
+		// The successor execution's CAS already stored its own token; a
+		// delayed mirror from the previous execution must not land.
+		seedResumeToken(t, repo, "s1", "t1", "acp-successor")
+		require.NoError(t, repo.SetSessionMetadataKey(ctx, "s1", "acp",
+			map[string]interface{}{"session_id": "acp-successor"}))
+		svc := createTestService(repo, newMockStepGetter(), newMockTaskRepo())
+
+		svc.persistACPSessionID(ctx, "s1", "acp-stale")
+
+		session, err := repo.GetTaskSession(ctx, "s1")
+		require.NoError(t, err)
+		acp := session.Metadata["acp"].(map[string]interface{})
+		require.Equal(t, "acp-successor", acp["session_id"], "stale id must not overwrite the live execution's id")
+	})
+}
@@ -165,6 +165,7 @@ type sessionExecutorStore interface {
 	UpdateSessionReviewStatus(ctx context.Context, sessionID string, status string) error
 	UpdateSessionMetadata(ctx context.Context, sessionID string, metadata map[string]interface{}) error
 	SetSessionMetadataKey(ctx context.Context, sessionID, key string, value interface{}) error
+	SetSessionACPSessionID(ctx context.Context, sessionID, acpSessionID string) (bool, error)
 	// Executor running state
 	ListExecutorsRunning(ctx context.Context) ([]*models.ExecutorRunning, error)
 	UpsertExecutorRunning(ctx context.Context, running *models.ExecutorRunning) error
@@ -535,6 +535,9 @@ func (m *mockRepository) SetSessionMetadataKey(ctx context.Context, sessionID, k
 	}
 	return nil
 }
+func (m *mockRepository) SetSessionACPSessionID(_ context.Context, _ string, _ string) (bool, error) {
+	return false, nil
+}
 func (m *mockRepository) DismissLastAgentError(_ context.Context, _ string, _ models.LastAgentError, _ time.Time) (bool, error) {
 	return true, nil
 }
@@ -184,6 +184,7 @@ type SessionRepository interface {
 	UpdateSessionReviewStatus(ctx context.Context, sessionID string, status string) error
 	UpdateSessionMetadata(ctx context.Context, sessionID string, metadata map[string]interface{}) error
 	SetSessionMetadataKey(ctx context.Context, sessionID, key string, value interface{}) error
+	SetSessionACPSessionID(ctx context.Context, sessionID, acpSessionID string) (bool, error)
 	DismissLastAgentError(ctx context.Context, sessionID string, expected models.LastAgentError, dismissedAt time.Time) (bool, error)
 	GetLastAgentMessage(ctx context.Context, sessionID string) (string, error)
 }
@@ -646,6 +646,38 @@ func (r *Repository) SetSessionMetadataKey(ctx context.Context, sessionID, key s
 	return nil
 }

+// SetSessionACPSessionID mirrors the agent's ACP session id into the session's
+// "acp" metadata map as a single atomic UPDATE. json_patch merges the sub-key,
+// so keys session_info already wrote (title, ...) survive without a
+// read-modify-write round trip. The write only happens while the session's
+// executors_running row still holds acpSessionID as its resume token — i.e.
+// the CAS that stored the token hasn't been superseded by a rotated execution
+// — and is skipped when the stored id is already current, so a no-op never
+// touches updated_at. Returns whether a row was written.
+func (r *Repository) SetSessionACPSessionID(ctx context.Context, sessionID, acpSessionID string) (bool, error) {
+	patch, err := json.Marshal(map[string]interface{}{"acp": map[string]string{"session_id": acpSessionID}})
+	if err != nil {
+		return false, fmt.Errorf("failed to serialize metadata patch: %w", err)
+	}
+	now := time.Now().UTC()
+	result, err := r.db.ExecContext(ctx, r.db.Rebind(`
+		UPDATE task_sessions
+		SET metadata = json_patch(CASE WHEN metadata IS NULL OR metadata = 'null' OR metadata = '' THEN '{}' ELSE metadata END, json(?)),
+		    updated_at = ?
+		WHERE id = ?
+		  AND json_extract(metadata, '$.acp.session_id') IS NOT ?
+		  AND EXISTS (
+			SELECT 1 FROM executors_running er
+			WHERE er.session_id = task_sessions.id AND er.resume_token = ?
+		  )
+	`), string(patch), now, sessionID, acpSessionID, acpSessionID)
+	if err != nil {
+		return false, err
+	}
+	rows, _ := result.RowsAffected()
+	return rows > 0, nil
+}
+
 func (r *Repository) DismissLastAgentError(ctx context.Context, sessionID string, expected models.LastAgentError, dismissedAt time.Time) (bool, error) {
 	next := expected
 	next.DismissedAt = &dismissedAt