feat: persist generic session attributes

pkg/runtime/agent_delegation.go

@@ -243,6 +243,7 @@ func newSubSession(parent *session.Session, cfg SubSessionConfig, childAgent *ag
 		session.WithSendUserMessage(false),
 		session.WithParentID(parent.ID),
 		session.WithAttachedFiles(attachedFiles),
+		session.WithAttributes(parent.AttributesSnapshot()),
 	}
 	if cfg.PinAgent {
 		opts = append(opts, session.WithAgentName(cfg.AgentName))

pkg/server/session_manager.go

@@ -567,6 +567,9 @@ func (sm *SessionManager) CreateSession(ctx context.Context, sessionTemplate *se
 	if sessionTemplate.Permissions != nil {
 		opts = append(opts, session.WithPermissions(sessionTemplate.Permissions))
 	}
+	if attributes := sessionTemplate.AttributesSnapshot(); len(attributes) > 0 {
+		opts = append(opts, session.WithAttributes(attributes))
+	}
 
 	sess := session.New(opts...)
 
@@ -1904,6 +1907,7 @@ func (sm *SessionManager) SetSessionAgentModel(ctx context.Context, sessionID, m
 		SafetyPolicy:            sess.SafetyPolicy,
 		ToolsApproved:           sess.ToolsApproved,
 		Permissions:             sess.Permissions,
+		Attributes:              sess.AttributesSnapshot(),
 		MaxIterations:           sess.MaxIterations,
 		MaxConsecutiveToolCalls: sess.MaxConsecutiveToolCalls,
 		MaxOldToolCallTokens:    sess.MaxOldToolCallTokens,

pkg/session/branch.go

@@ -100,6 +100,7 @@ func (s *Session) Clone() *Session {
 		OutputTokens:            s.OutputTokens,
 		Cost:                    s.Cost,
 		Permissions:             s.Permissions.Clone(),
+		Attributes:              maps.Clone(s.Attributes),
 		AgentModelOverrides:     cloneStringMap(s.AgentModelOverrides),
 		CustomModelsUsed:        cloneStringSlice(s.CustomModelsUsed),
 		AttachedFiles:           cloneStringSlice(s.AttachedFiles),
@@ -228,6 +229,7 @@ func copySessionMetadata(dst, src *Session, title string) {
 	dst.MaxToolResultTokens = src.MaxToolResultTokens
 	dst.Starred = src.Starred
 	dst.Permissions = src.Permissions.Clone()
+	dst.Attributes = src.AttributesSnapshot()
 	dst.AgentModelOverrides = cloneStringMap(src.AgentModelOverrides)
 	dst.CustomModelsUsed = cloneStringSlice(src.CustomModelsUsed)
 	dst.AttachedFiles = src.AttachedFilesSnapshot()

pkg/session/migrations.go

@@ -428,6 +428,13 @@ func getAllMigrations() []Migration {
 				ALTER TABLE session_items ADD COLUMN usage_json TEXT NOT NULL DEFAULT '';
 			`,
 		},
+		{
+			ID:          26,
+			Name:        "026_add_session_attributes_column",
+			Description: "Add generic attributes to sessions",
+			UpSQL:       `ALTER TABLE sessions ADD COLUMN attributes TEXT DEFAULT '{}'`,
+			DownSQL:     `ALTER TABLE sessions DROP COLUMN attributes`,
+		},
 	}
 }
 

pkg/session/migrations_pinned_test.go

@@ -39,7 +39,7 @@ func TestMigrationCatalogIsContentPinned(t *testing.T) {
 
 	got := digestMigrationCatalog(getAllMigrations())
 
-	const wantDigest = "fa29f858ecfe989a2247769056048116c6f9a566a91057fd28c9b2f3b964d66b"
+	const wantDigest = "71e0d68cf3a8439361a8339bad4ac543ffe3fda02fcce480ca0adc4e0f4017ce"
 	if got != wantDigest {
 		t.Fatalf(`migration catalogue content has changed.
 

pkg/session/session.go

@@ -211,10 +211,9 @@ type Termination struct {
 
 // Session represents the agent's state including conversation history and variables
 type Session struct {
-	// mu protects Messages and the scalar metadata that is written
-	// cross-goroutine (Title, InputTokens, OutputTokens, Cost, ...) from
-	// concurrent read/write access. Shared-session readers must go through
-	// the locked accessors (TitleSnapshot, Usage, TokensAndCost, ...).
+	// mu protects Messages and metadata that is written cross-goroutine
+	// (Title, Attributes, InputTokens, OutputTokens, Cost, ...) from concurrent
+	// read/write access. Shared-session readers must use the locked accessors.
 	mu sync.RWMutex `json:"-"`
 
 	// now and newID are per-session sources of time and identity. They are
@@ -323,6 +322,11 @@ type Session struct {
 	// When set, these are evaluated before team-level permissions.
 	Permissions *PermissionsConfig `json:"permissions,omitempty"`
 
+	// Attributes stores generic, namespaced metadata supplied by embedders.
+	// Shared-session callers must use AttributesSnapshot, SetAttribute, and
+	// DeleteAttribute rather than mutating this map directly.
+	Attributes map[string]string `json:"attributes,omitempty"`
+
 	// AgentModelOverrides stores per-agent model overrides for this session.
 	// Key is the agent name, value is the model reference (e.g., "openai/gpt-4o" or a named model from config).
 	// When a session is loaded, these overrides are reapplied to the runtime.
@@ -786,6 +790,19 @@ func (s *Session) AddMessage(msg *Message) int {
 	return len(s.Messages) - 1
 }
 
+// MarshalJSON takes a consistent snapshot while encoding mutable session
+// state. In particular, SetAttribute may otherwise mutate a map while the JSON
+// encoder is iterating it.
+func (s *Session) MarshalJSON() ([]byte, error) {
+	if s == nil {
+		return []byte("null"), nil
+	}
+	s.mu.RLock()
+	defer s.mu.RUnlock()
+	type sessionJSON Session
+	return json.Marshal((*sessionJSON)(s))
+}
+
 // SetUsage records cumulative input/output token counts under s.mu.
 // The runtime stream goroutine and the persistence observer race on
 // these fields without it.
@@ -1263,6 +1280,38 @@ func (s *Session) AttachedFilesSnapshot() []string {
 	return slices.Clone(s.AttachedFiles)
 }
 
+// AttributesSnapshot returns an independent copy of the session attributes.
+func (s *Session) AttributesSnapshot() map[string]string {
+	s.mu.RLock()
+	defer s.mu.RUnlock()
+	return maps.Clone(s.Attributes)
+}
+
+// SetAttribute sets a session attribute. Empty keys are ignored because they
+// cannot form a meaningful namespaced metadata key.
+func (s *Session) SetAttribute(key, value string) {
+	if key == "" {
+		return
+	}
+	s.mu.Lock()
+	defer s.mu.Unlock()
+	if s.Attributes == nil {
+		s.Attributes = make(map[string]string)
+	}
+	s.Attributes[key] = value
+}
+
+// DeleteAttribute deletes a session attribute. An empty key is a no-op,
+// matching SetAttribute and WithAttributes.
+func (s *Session) DeleteAttribute(key string) {
+	if key == "" {
+		return
+	}
+	s.mu.Lock()
+	defer s.mu.Unlock()
+	delete(s.Attributes, key)
+}
+
 // DelegationLineageSnapshot returns a copy of the session's delegation
 // lineage. Callers may freely mutate the returned slice without affecting
 // the session.
@@ -1274,6 +1323,18 @@ func (s *Session) DelegationLineageSnapshot() []string {
 
 type Opt func(s *Session)
 
+// WithAttributes sets generic session metadata from an independent copy of
+// attributes. Empty keys are discarded; empty values are preserved.
+func WithAttributes(attributes map[string]string) Opt {
+	cloned := maps.Clone(attributes)
+	delete(cloned, "")
+	return func(s *Session) {
+		s.mu.Lock()
+		defer s.mu.Unlock()
+		s.Attributes = maps.Clone(cloned)
+	}
+}
+
 func WithUserMessage(content string) Opt {
 	return func(s *Session) {
 		s.AddMessage(UserMessageAt(s.now(), content))

pkg/session/session_attributes_test.go

@@ -0,0 +1,92 @@
+package session
+
+import (
+	"encoding/json"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+func TestWithAttributesClonesInput(t *testing.T) {
+	t.Parallel()
+	input := map[string]string{
+		"daw.workspace_path": "/workspace",
+		"":                   "ignored",
+	}
+
+	sess := New(WithAttributes(input))
+	input["daw.workspace_path"] = "/mutated"
+	input["daw.worktree_id"] = "new"
+
+	assert.Equal(t, map[string]string{"daw.workspace_path": "/workspace"}, sess.AttributesSnapshot())
+}
+
+func TestAttributesSnapshotClonesOutput(t *testing.T) {
+	t.Parallel()
+	sess := New(WithAttributes(map[string]string{"daw.worktree_id": "one"}))
+
+	snapshot := sess.AttributesSnapshot()
+	snapshot["daw.worktree_id"] = "two"
+	snapshot["daw.worktree_path"] = "/other"
+
+	assert.Equal(t, map[string]string{"daw.worktree_id": "one"}, sess.AttributesSnapshot())
+}
+
+func TestSetAndDeleteAttribute(t *testing.T) {
+	t.Parallel()
+	sess := New()
+
+	sess.SetAttribute("daw.execution_type", "worktree")
+	sess.SetAttribute("daw.worktree_id", "one")
+	sess.SetAttribute("", "ignored")
+	assert.Equal(t, map[string]string{
+		"daw.execution_type": "worktree",
+		"daw.worktree_id":    "one",
+	}, sess.AttributesSnapshot())
+
+	sess.SetAttribute("daw.worktree_id", "two")
+	sess.DeleteAttribute("daw.execution_type")
+	sess.DeleteAttribute("")
+	assert.Equal(t, map[string]string{"daw.worktree_id": "two"}, sess.AttributesSnapshot())
+}
+
+func TestSessionAttributesJSONRoundTrip(t *testing.T) {
+	t.Parallel()
+	original := New(WithAttributes(map[string]string{
+		"daw.workspace_path": "/workspace",
+		"daw.worktree_id":    "wt-1",
+	}))
+
+	data, err := json.Marshal(original)
+	require.NoError(t, err)
+	assert.Contains(t, string(data), `"attributes"`)
+
+	var decoded Session
+	require.NoError(t, json.Unmarshal(data, &decoded))
+	assert.Equal(t, original.AttributesSnapshot(), decoded.AttributesSnapshot())
+
+	decoded.SetAttribute("daw.worktree_id", "wt-2")
+	assert.Equal(t, "wt-1", original.AttributesSnapshot()["daw.worktree_id"])
+}
+
+func TestCloneAndBranchCopyAttributesIndependently(t *testing.T) {
+	t.Parallel()
+	parent := New(
+		WithAttributes(map[string]string{"daw.worktree_id": "wt-1"}),
+		WithUserMessage("hello"),
+	)
+
+	clone := parent.Clone()
+	branched, err := BranchSession(parent, 1)
+	require.NoError(t, err)
+
+	assert.Equal(t, parent.AttributesSnapshot(), clone.AttributesSnapshot())
+	assert.Equal(t, parent.AttributesSnapshot(), branched.AttributesSnapshot())
+
+	clone.SetAttribute("daw.worktree_id", "clone")
+	branched.SetAttribute("daw.worktree_id", "branch")
+	assert.Equal(t, "wt-1", parent.AttributesSnapshot()["daw.worktree_id"])
+	assert.Equal(t, "clone", clone.AttributesSnapshot()["daw.worktree_id"])
+	assert.Equal(t, "branch", branched.AttributesSnapshot()["daw.worktree_id"])
+}

pkg/session/session_race_test.go

@@ -1,12 +1,37 @@
 package session
 
 import (
+	"encoding/json"
 	"sync"
 	"testing"
 
 	"github.com/docker/docker-agent/pkg/chat"
 )
 
+func TestSessionAttributesConcurrent(t *testing.T) {
+	t.Parallel()
+
+	s := New(WithAttributes(map[string]string{"daw.workspace_path": "/workspace"}))
+	var wg sync.WaitGroup
+	for range 100 {
+		wg.Go(func() {
+			s.SetAttribute("daw.worktree_id", "worktree")
+		})
+		wg.Go(func() {
+			_ = s.AttributesSnapshot()
+		})
+		wg.Go(func() {
+			s.DeleteAttribute("daw.worktree_id")
+		})
+		wg.Go(func() {
+			if _, err := json.Marshal(s); err != nil {
+				t.Errorf("Marshal: %v", err)
+			}
+		})
+	}
+	wg.Wait()
+}
+
 func TestAddMessageUsageRecordConcurrent(t *testing.T) {
 	t.Parallel()