diff --git a/pkg/runtime/agent_delegation.go b/pkg/runtime/agent_delegation.go index d595185c9e..265eef498c 100644 --- a/pkg/runtime/agent_delegation.go +++ b/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)) diff --git a/pkg/server/session_manager.go b/pkg/server/session_manager.go index 5fe834f15e..0c9896d8b1 100644 --- a/pkg/server/session_manager.go +++ b/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, diff --git a/pkg/session/branch.go b/pkg/session/branch.go index 5ee442b86e..ac639eaeac 100644 --- a/pkg/session/branch.go +++ b/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() diff --git a/pkg/session/migrations.go b/pkg/session/migrations.go index 7a75f06211..8534d1c4a8 100644 --- a/pkg/session/migrations.go +++ b/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`, + }, } } diff --git a/pkg/session/migrations_pinned_test.go b/pkg/session/migrations_pinned_test.go index 755f26b89f..e4dc5a45ad 100644 --- a/pkg/session/migrations_pinned_test.go +++ b/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. diff --git a/pkg/session/session.go b/pkg/session/session.go index 2055385019..09804298ed 100644 --- a/pkg/session/session.go +++ b/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)) diff --git a/pkg/session/session_attributes_test.go b/pkg/session/session_attributes_test.go new file mode 100644 index 0000000000..6b89dcffd8 --- /dev/null +++ b/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"]) +} diff --git a/pkg/session/session_race_test.go b/pkg/session/session_race_test.go index cfcdd6d13f..3f53cad67f 100644 --- a/pkg/session/session_race_test.go +++ b/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() diff --git a/pkg/session/store.go b/pkg/session/store.go index 3a7321bf4b..7ac10175ab 100644 --- a/pkg/session/store.go +++ b/pkg/session/store.go @@ -7,6 +7,7 @@ import ( "errors" "fmt" "log/slog" + "maps" "slices" "strconv" "strings" @@ -82,6 +83,7 @@ type Summary struct { Starred bool NumMessages int WorkingDir string + Attributes map[string]string } // Store defines the interface for session storage @@ -185,6 +187,7 @@ func (s *InMemorySessionStore) GetSessionSummaries(_ context.Context) ([]Summary Starred: value.Starred, NumMessages: value.MessageCount(), WorkingDir: value.WorkingDir, + Attributes: value.AttributesSnapshot(), }) return true }) @@ -235,6 +238,7 @@ func (s *InMemorySessionStore) UpdateSession(_ context.Context, session *Session OutputTokens: session.OutputTokens, Cost: session.Cost, Permissions: session.Permissions.Clone(), + Attributes: maps.Clone(session.Attributes), AgentModelOverrides: cloneStringMap(session.AgentModelOverrides), CustomModelsUsed: cloneStringSlice(session.CustomModelsUsed), InstructionContext: cloneInstructionContext(session.InstructionContext), @@ -376,7 +380,7 @@ type SQLiteSessionStore struct { // sessionSelectColumns is the canonical SELECT list for the sessions table. // The column order matches what scanSession expects; all read paths use this // constant so that adding a column requires updating exactly one place. -const sessionSelectColumns = `id, tools_approved, safety_policy, input_tokens, output_tokens, title, cost, send_user_message, max_iterations, working_dir, created_at, starred, permissions, agent_model_overrides, custom_models_used, thinking, parent_id, instruction_context` +const sessionSelectColumns = `id, tools_approved, safety_policy, input_tokens, output_tokens, title, cost, send_user_message, max_iterations, working_dir, created_at, starred, permissions, agent_model_overrides, custom_models_used, thinking, parent_id, instruction_context, attributes` // sessionPersistedFields holds the encoded form of a Session's JSON-bearing // columns plus the SQL representation of parent_id (nil for the empty @@ -386,6 +390,7 @@ type sessionPersistedFields struct { AgentModelOverridesJSON string CustomModelsUsedJSON string InstructionContextJSON string + AttributesJSON string ParentID any // string or nil } @@ -396,6 +401,16 @@ type sessionPersistedFields struct { func sessionPersistedFieldsOf(session *Session) (sessionPersistedFields, error) { var f sessionPersistedFields + attributes := session.AttributesSnapshot() + f.AttributesJSON = "{}" + if len(attributes) > 0 { + attributesBytes, err := json.Marshal(attributes) + if err != nil { + return f, err + } + f.AttributesJSON = string(attributesBytes) + } + if session.Permissions != nil { permBytes, err := json.Marshal(session.Permissions) if err != nil { @@ -541,12 +556,12 @@ func (s *SQLiteSessionStore) AddSession(ctx context.Context, session *Session) e `INSERT INTO sessions ( id, tools_approved, safety_policy, input_tokens, output_tokens, title, cost, send_user_message, max_iterations, working_dir, created_at, permissions, agent_model_overrides, - custom_models_used, thinking, parent_id, instruction_context - ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)`, + custom_models_used, thinking, parent_id, instruction_context, attributes + ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)`, session.ID, session.ToolsApproved, string(session.SafetyPolicy), session.InputTokens, session.OutputTokens, session.Title, session.Cost, session.SendUserMessage, session.MaxIterations, session.WorkingDir, session.CreatedAt.Format(time.RFC3339), fields.PermissionsJSON, fields.AgentModelOverridesJSON, - fields.CustomModelsUsedJSON, false, fields.ParentID, fields.InstructionContextJSON) + fields.CustomModelsUsedJSON, false, fields.ParentID, fields.InstructionContextJSON, fields.AttributesJSON) if err != nil { return err } @@ -578,6 +593,7 @@ func scanSession(scanner interface { agentModelOverridesJSON string customModelsUsedJSON string instructionContextJSON sql.NullString + attributesJSON sql.NullString createdAtStr string thinking bool // discarded ) @@ -586,7 +602,7 @@ func scanSession(scanner interface { &sess.ID, &sess.ToolsApproved, &safetyPolicy, &sess.InputTokens, &sess.OutputTokens, &sess.Title, &sess.Cost, &sess.SendUserMessage, &sess.MaxIterations, &workingDir, &createdAtStr, &sess.Starred, &permissionsJSON, - &agentModelOverridesJSON, &customModelsUsedJSON, &thinking, &parentID, &instructionContextJSON, + &agentModelOverridesJSON, &customModelsUsedJSON, &thinking, &parentID, &instructionContextJSON, &attributesJSON, ) if err != nil { return nil, err @@ -623,9 +639,28 @@ func scanSession(scanner interface { } } + sess.Attributes, err = decodeAttributes(attributesJSON) + if err != nil { + return nil, fmt.Errorf("unmarshaling session attributes: %w", err) + } + return &sess, nil } +func decodeAttributes(value sql.NullString) (map[string]string, error) { + if !value.Valid || strings.TrimSpace(value.String) == "" { + return nil, nil + } + var attributes map[string]string + if err := json.Unmarshal([]byte(value.String), &attributes); err != nil { + return nil, err + } + if len(attributes) == 0 { + return nil, nil + } + return attributes, nil +} + // GetSession retrieves a session by ID func (s *SQLiteSessionStore) GetSession(ctx context.Context, id string) (*Session, error) { if id == "" { @@ -808,7 +843,7 @@ func (s *SQLiteSessionStore) GetSessions(ctx context.Context) ([]*Session, error // This is much faster than GetSessions as it doesn't load message content. func (s *SQLiteSessionStore) GetSessionSummaries(ctx context.Context) ([]Summary, error) { rows, err := s.db.QueryContext(ctx, - `SELECT s.id, s.title, s.created_at, s.starred, s.working_dir, + `SELECT s.id, s.title, s.created_at, s.starred, s.working_dir, s.attributes, (SELECT COUNT(*) FROM session_items si WHERE si.session_id = s.id AND si.item_type = 'message') FROM sessions s WHERE s.parent_id IS NULL OR s.parent_id = '' @@ -821,15 +856,20 @@ func (s *SQLiteSessionStore) GetSessionSummaries(ctx context.Context) ([]Summary var summaries []Summary for rows.Next() { var ( - summary Summary - createdAtStr string - workingDir sql.NullString + summary Summary + createdAtStr string + workingDir sql.NullString + attributesJSON sql.NullString ) - if err := rows.Scan(&summary.ID, &summary.Title, &createdAtStr, &summary.Starred, &workingDir, &summary.NumMessages); err != nil { + if err := rows.Scan(&summary.ID, &summary.Title, &createdAtStr, &summary.Starred, &workingDir, &attributesJSON, &summary.NumMessages); err != nil { return nil, err } summary.CreatedAt = parseCreatedAt(createdAtStr) summary.WorkingDir = workingDir.String + summary.Attributes, err = decodeAttributes(attributesJSON) + if err != nil { + return nil, fmt.Errorf("unmarshaling session attributes for %s: %w", summary.ID, err) + } summaries = append(summaries, summary) } @@ -892,6 +932,7 @@ func (s *SQLiteSessionStore) UpdateSession(ctx context.Context, session *Session OutputTokens: session.OutputTokens, Cost: session.Cost, Permissions: session.Permissions.Clone(), + Attributes: maps.Clone(session.Attributes), AgentModelOverrides: cloneStringMap(session.AgentModelOverrides), CustomModelsUsed: cloneStringSlice(session.CustomModelsUsed), InstructionContext: cloneInstructionContext(session.InstructionContext), @@ -916,9 +957,9 @@ func (s *SQLiteSessionStore) UpdateSession(ctx context.Context, session *Session `INSERT INTO sessions ( id, tools_approved, safety_policy, input_tokens, output_tokens, title, cost, send_user_message, max_iterations, working_dir, created_at, starred, permissions, agent_model_overrides, - custom_models_used, thinking, parent_id, instruction_context + custom_models_used, thinking, parent_id, instruction_context, attributes ) - VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) + VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) ON CONFLICT(id) DO UPDATE SET title = excluded.title, tools_approved = excluded.tools_approved, @@ -935,11 +976,12 @@ func (s *SQLiteSessionStore) UpdateSession(ctx context.Context, session *Session custom_models_used = excluded.custom_models_used, thinking = excluded.thinking, parent_id = excluded.parent_id, - instruction_context = excluded.instruction_context`, + instruction_context = excluded.instruction_context, + attributes = excluded.attributes`, snapshot.ID, snapshot.ToolsApproved, string(snapshot.SafetyPolicy), snapshot.InputTokens, snapshot.OutputTokens, snapshot.Title, snapshot.Cost, snapshot.SendUserMessage, snapshot.MaxIterations, snapshot.WorkingDir, snapshot.CreatedAt.Format(time.RFC3339), snapshot.Starred, fields.PermissionsJSON, fields.AgentModelOverridesJSON, - fields.CustomModelsUsedJSON, false, fields.ParentID, fields.InstructionContextJSON) + fields.CustomModelsUsedJSON, false, fields.ParentID, fields.InstructionContextJSON, fields.AttributesJSON) if err != nil { return err } @@ -1086,14 +1128,14 @@ func (s *SQLiteSessionStore) addSessionTx(ctx context.Context, tx *sql.Tx, sessi `INSERT INTO sessions ( id, tools_approved, safety_policy, input_tokens, output_tokens, title, cost, send_user_message, max_iterations, working_dir, created_at, starred, permissions, agent_model_overrides, - custom_models_used, thinking, parent_id, instruction_context + custom_models_used, thinking, parent_id, instruction_context, attributes ) - VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)`, + VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)`, session.ID, session.ToolsApproved, string(session.SafetyPolicy), session.InputTokens, session.OutputTokens, session.Title, session.Cost, session.SendUserMessage, session.MaxIterations, session.WorkingDir, session.CreatedAt.Format(time.RFC3339), session.Starred, fields.PermissionsJSON, fields.AgentModelOverridesJSON, fields.CustomModelsUsedJSON, false, - fields.ParentID, fields.InstructionContextJSON) + fields.ParentID, fields.InstructionContextJSON, fields.AttributesJSON) return err } diff --git a/pkg/session/store_attributes_test.go b/pkg/session/store_attributes_test.go new file mode 100644 index 0000000000..c313a31114 --- /dev/null +++ b/pkg/session/store_attributes_test.go @@ -0,0 +1,137 @@ +package session + +import ( + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func TestSQLiteSessionAttributesAddRoundTrip(t *testing.T) { + t.Parallel() + store := openMemoryStore(t) + sess := New( + WithID("attributes-add"), + WithAttributes(map[string]string{ + "daw.workspace_path": "/workspace", + "daw.worktree_id": "wt-1", + }), + ) + + require.NoError(t, store.AddSession(t.Context(), sess)) + loaded, err := store.GetSession(t.Context(), sess.ID) + require.NoError(t, err) + assert.Equal(t, sess.AttributesSnapshot(), loaded.AttributesSnapshot()) +} + +func TestSQLiteSessionAttributesUpdateRoundTrip(t *testing.T) { + t.Parallel() + store := openMemoryStore(t) + sess := New(WithID("attributes-update"), WithAttributes(map[string]string{"daw.worktree_id": "wt-1"})) + require.NoError(t, store.AddSession(t.Context(), sess)) + + sess.SetAttribute("daw.worktree_id", "wt-2") + sess.SetAttribute("daw.worktree_branch", "feature") + require.NoError(t, store.UpdateSession(t.Context(), sess)) + + loaded, err := store.GetSession(t.Context(), sess.ID) + require.NoError(t, err) + assert.Equal(t, map[string]string{ + "daw.worktree_id": "wt-2", + "daw.worktree_branch": "feature", + }, loaded.AttributesSnapshot()) +} + +func TestSQLiteUnrelatedMetadataUpdatePreservesAttributes(t *testing.T) { + t.Parallel() + store := openMemoryStore(t) + sess := New(WithID("attributes-preserved"), WithAttributes(map[string]string{"daw.execution_type": "worktree"})) + require.NoError(t, store.AddSession(t.Context(), sess)) + + loaded, err := store.GetSession(t.Context(), sess.ID) + require.NoError(t, err) + loaded.SetTitle("renamed") + require.NoError(t, store.UpdateSession(t.Context(), loaded)) + require.NoError(t, store.UpdateSessionTitle(t.Context(), sess.ID, "renamed again")) + + reloaded, err := store.GetSession(t.Context(), sess.ID) + require.NoError(t, err) + assert.Equal(t, map[string]string{"daw.execution_type": "worktree"}, reloaded.AttributesSnapshot()) +} + +func TestSQLiteSessionSummariesIncludeAttributesWithoutLoadingMessages(t *testing.T) { + t.Parallel() + store := openMemoryStore(t) + sess := New( + WithID("attributes-summary"), + WithAttributes(map[string]string{"daw.workspace_path": "/workspace"}), + WithUserMessage("hello"), + ) + require.NoError(t, store.AddSession(t.Context(), sess)) + + // A full load would fail on this malformed payload. Summaries must obtain + // attributes directly from the sessions row without loading item history. + _, err := store.db.ExecContext(t.Context(), + "UPDATE session_items SET message_json = 'not-json' WHERE session_id = ?", sess.ID) + require.NoError(t, err) + + summaries, err := store.GetSessionSummaries(t.Context()) + require.NoError(t, err) + require.Len(t, summaries, 1) + assert.Equal(t, map[string]string{"daw.workspace_path": "/workspace"}, summaries[0].Attributes) + + summaries[0].Attributes["daw.workspace_path"] = "/mutated" + assert.Equal(t, "/workspace", sess.AttributesSnapshot()["daw.workspace_path"]) +} + +func TestInMemorySessionSummaryAttributesAreIndependent(t *testing.T) { + t.Parallel() + store := NewInMemorySessionStore() + sess := New(WithID("memory-attributes"), WithAttributes(map[string]string{"daw.worktree_id": "wt-1"})) + require.NoError(t, store.AddSession(t.Context(), sess)) + + summaries, err := store.GetSessionSummaries(t.Context()) + require.NoError(t, err) + require.Len(t, summaries, 1) + summaries[0].Attributes["daw.worktree_id"] = "mutated" + summaries[0].Attributes["daw.worktree_path"] = "/other" + + assert.Equal(t, map[string]string{"daw.worktree_id": "wt-1"}, sess.AttributesSnapshot()) +} + +func TestInMemorySessionAttributesUpdateRoundTrip(t *testing.T) { + t.Parallel() + store := NewInMemorySessionStore() + sess := New(WithID("memory-update"), WithAttributes(map[string]string{"daw.worktree_id": "wt-1"})) + require.NoError(t, store.UpdateSession(t.Context(), sess)) + + sess.SetAttribute("daw.worktree_id", "mutated-after-update") + loaded, err := store.GetSession(t.Context(), sess.ID) + require.NoError(t, err) + assert.Equal(t, map[string]string{"daw.worktree_id": "wt-1"}, loaded.AttributesSnapshot()) +} + +func TestSQLiteSessionAttributesLegacyNullAndEmptyValues(t *testing.T) { + t.Parallel() + store := openMemoryStore(t) + for _, id := range []string{"attributes-null", "attributes-empty"} { + require.NoError(t, store.AddSession(t.Context(), New(WithID(id)))) + } + _, err := store.db.ExecContext(t.Context(), "UPDATE sessions SET attributes = NULL WHERE id = 'attributes-null'") + require.NoError(t, err) + _, err = store.db.ExecContext(t.Context(), "UPDATE sessions SET attributes = '' WHERE id = 'attributes-empty'") + require.NoError(t, err) + + for _, id := range []string{"attributes-null", "attributes-empty"} { + loaded, err := store.GetSession(t.Context(), id) + require.NoError(t, err) + assert.Empty(t, loaded.AttributesSnapshot()) + } + + summaries, err := store.GetSessionSummaries(t.Context()) + require.NoError(t, err) + require.Len(t, summaries, 2) + for _, summary := range summaries { + assert.Empty(t, summary.Attributes) + } +} diff --git a/pkg/session/store_persisted_fields_test.go b/pkg/session/store_persisted_fields_test.go index 07909a9f50..ad548b4ad9 100644 --- a/pkg/session/store_persisted_fields_test.go +++ b/pkg/session/store_persisted_fields_test.go @@ -13,6 +13,7 @@ func TestSessionPersistedFieldsOf_Defaults(t *testing.T) { require.NoError(t, err) assert.Empty(t, got.PermissionsJSON, "nil permissions should serialise to empty string") + assert.Equal(t, "{}", got.AttributesJSON, "nil attributes should default to {}") assert.Equal(t, "{}", got.AgentModelOverridesJSON, "nil overrides should default to {}") assert.Equal(t, "[]", got.CustomModelsUsedJSON, "nil models should default to []") assert.Nil(t, got.ParentID, "empty parent_id must encode as SQL NULL") @@ -26,6 +27,7 @@ func TestSessionPersistedFieldsOf_PopulatedValues(t *testing.T) { Permissions: &PermissionsConfig{ Allow: []string{"shell"}, }, + Attributes: map[string]string{"daw.workspace_path": "/workspace"}, AgentModelOverrides: map[string]string{"root": "openai/gpt-5"}, CustomModelsUsed: []string{"openai/gpt-5"}, } @@ -34,6 +36,7 @@ func TestSessionPersistedFieldsOf_PopulatedValues(t *testing.T) { require.NoError(t, err) assert.JSONEq(t, `{"allow":["shell"]}`, got.PermissionsJSON) + assert.JSONEq(t, `{"daw.workspace_path":"/workspace"}`, got.AttributesJSON) assert.JSONEq(t, `{"root":"openai/gpt-5"}`, got.AgentModelOverridesJSON) assert.JSONEq(t, `["openai/gpt-5"]`, got.CustomModelsUsedJSON) assert.Equal(t, "parent-1", got.ParentID) @@ -44,10 +47,12 @@ func TestSessionPersistedFieldsOf_EmptyMapsAndSlicesUseDefaults(t *testing.T) { // len() == 0 for non-nil empty values must take the default branch // because JSON encoding "{}" / "[]" matches what the schema expects. got, err := sessionPersistedFieldsOf(&Session{ + Attributes: map[string]string{}, AgentModelOverrides: map[string]string{}, CustomModelsUsed: []string{}, }) require.NoError(t, err) + assert.Equal(t, "{}", got.AttributesJSON) assert.Equal(t, "{}", got.AgentModelOverridesJSON) assert.Equal(t, "[]", got.CustomModelsUsedJSON) }