From 458b7808ee8c020970398951b282826299fed77b Mon Sep 17 00:00:00 2001 From: Trung Nguyen Date: Thu, 16 Jul 2026 15:40:26 +0200 Subject: [PATCH 1/2] fix(session): persist SafetyPolicy across turns The Session struct grew a SafetyPolicy field, and the API + dispatcher learned to read and write it, but none of the store code carried it through. As a result: - The SQLite schema had no column for it, so every UpdateSession silently dropped the value on write and every scanSession returned the empty string on read. - The InMemory and SQLite UpdateSession snapshot builders copy fields explicitly for lock-safety; both were missing SafetyPolicy, so even once the column existed the value would still be dropped. - SessionManager.ResumeSession forwarded approve-safe / approve-safer / approve-session to the runtime and returned. The dispatcher mutates the in-memory session pointer to record the new policy, but PersistenceObserver only calls UpdateSession on OnRunStart. Nothing else re-persisted the row before the next turn. The net effect: a caller could opt into safe-auto (or any of the other mid-turn policy transitions), the runtime would honor it for the rest of the current turn, and then the next turn would reload the session from storage and see an empty SafetyPolicy, prompting again for everything the opt-in was supposed to skip. Fix: - Add a migration for a safety_policy column and thread it through scanSession, sessionSelectColumns, and every INSERT / ON CONFLICT UPDATE that persists a session row. - Include SafetyPolicy in both UpdateSession snapshot builders. - In SessionManager.ResumeSession, apply the same mutation the dispatcher will apply and persist synchronously before forwarding. - In SetSessionSafetyPolicy, mirror onto the live runtime session when one is attached so the current turn also picks up the change. --- pkg/server/session_manager.go | 36 ++++++++++++++++++++++++++- pkg/session/migrations.go | 6 +++++ pkg/session/migrations_pinned_test.go | 2 +- pkg/session/store.go | 27 ++++++++++++-------- 4 files changed, 58 insertions(+), 13 deletions(-) diff --git a/pkg/server/session_manager.go b/pkg/server/session_manager.go index 7e0752bd7..d5a3369ef 100644 --- a/pkg/server/session_manager.go +++ b/pkg/server/session_manager.go @@ -914,12 +914,37 @@ func (sm *SessionManager) ResumeSession(ctx context.Context, sessionID, confirma sm.mux.Lock() defer sm.mux.Unlock() - // Ensure the session runtime exists rt, exists := sm.runtimeSessions.Load(sessionID) if !exists { return errors.New("session not found") } + // approve-safe / approve-safer / approve-session mutate the session + // mid-turn from the dispatcher, but PersistenceObserver only persists + // on OnRunStart — the mutation would otherwise not survive to the + // next turn. Apply and persist synchronously here so the state is + // durable regardless of when the dispatcher processes the resume. + if rt.session != nil { + mutated := false + switch runtime.ResumeType(confirmation) { + case runtime.ResumeTypeApproveSafe: + rt.session.SetSafetyPolicy(session.SafetyPolicySafeAuto) + mutated = true + case runtime.ResumeTypeApproveSafer: + rt.session.SetSafetyPolicy(session.SafetyPolicySafer) + mutated = true + case runtime.ResumeTypeApproveSession: + rt.session.SetToolsApproved(true) + mutated = true + } + if mutated { + if err := sm.sessionStore.UpdateSession(ctx, rt.session); err != nil { + slog.WarnContext(ctx, "failed to persist mid-turn session state", + "session_id", sessionID, "confirmation", confirmation, "err", err) + } + } + } + rt.runtime.Resume(ctx, runtime.ResumeRequest{ Type: runtime.ResumeType(confirmation), Reason: reason, @@ -1130,6 +1155,15 @@ func (sm *SessionManager) SetSessionSafetyPolicy(ctx context.Context, sessionID } sm.mux.Lock() defer sm.mux.Unlock() + + // Mirror onto the live runtime session (if any) so the running + // dispatcher picks up the new policy on the very next tool call, + // then persist so it survives to the next turn. + if rt, ok := sm.runtimeSessions.Load(sessionID); ok && rt.session != nil { + rt.session.SetSafetyPolicy(policy) + return sm.sessionStore.UpdateSession(ctx, rt.session) + } + sess, err := sm.sessionStore.GetSession(ctx, sessionID) if err != nil { return err diff --git a/pkg/session/migrations.go b/pkg/session/migrations.go index 0d892a6be..c4f321faf 100644 --- a/pkg/session/migrations.go +++ b/pkg/session/migrations.go @@ -413,6 +413,12 @@ func getAllMigrations() []Migration { Description: "Add cost column to session_items so compaction summary costs survive reload", UpSQL: `ALTER TABLE session_items ADD COLUMN cost REAL NOT NULL DEFAULT 0`, }, + { + ID: 24, + Name: "024_add_safety_policy_column", + Description: "Add safety_policy column to sessions table", + UpSQL: `ALTER TABLE sessions ADD COLUMN safety_policy TEXT DEFAULT ''`, + }, } } diff --git a/pkg/session/migrations_pinned_test.go b/pkg/session/migrations_pinned_test.go index 1743c3bc3..65351363e 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 = "19cc1ae9c44d6ba716afea1f0f7f1a87623caf0f4540a751613e988ee2fd7049" + const wantDigest = "ba6212344ecacb721bda5a00a5f60ad02a891b79d8a12103484252f78b597569" if got != wantDigest { t.Fatalf(`migration catalogue content has changed. diff --git a/pkg/session/store.go b/pkg/session/store.go index 97a8a187e..8480c852c 100644 --- a/pkg/session/store.go +++ b/pkg/session/store.go @@ -226,6 +226,7 @@ func (s *InMemorySessionStore) UpdateSession(_ context.Context, session *Session Evals: session.Evals, CreatedAt: session.CreatedAt, ToolsApproved: session.ToolsApproved, + SafetyPolicy: session.SafetyPolicy, HideToolResults: session.HideToolResults, WorkingDir: session.WorkingDir, SendUserMessage: session.SendUserMessage, @@ -376,7 +377,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, 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` // sessionPersistedFields holds the encoded form of a Session's JSON-bearing // columns plus the SQL representation of parent_id (nil for the empty @@ -643,11 +644,11 @@ func (s *SQLiteSessionStore) AddSession(ctx context.Context, session *Session) e _, err = tx.ExecContext(ctx, `INSERT INTO sessions ( - id, tools_approved, input_tokens, output_tokens, title, cost, send_user_message, + 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 (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)`, - session.ID, session.ToolsApproved, session.InputTokens, session.OutputTokens, session.Title, + ) 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) @@ -675,6 +676,7 @@ func scanSession(scanner interface { ) (*Session, error) { var ( sess Session + safetyPolicy sql.NullString workingDir sql.NullString permissionsJSON sql.NullString parentID sql.NullString @@ -686,7 +688,7 @@ func scanSession(scanner interface { ) err := scanner.Scan( - &sess.ID, &sess.ToolsApproved, &sess.InputTokens, &sess.OutputTokens, + &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, @@ -696,6 +698,7 @@ func scanSession(scanner interface { } sess.CreatedAt = parseCreatedAt(createdAtStr) + sess.SafetyPolicy = SafetyPolicy(safetyPolicy.String) sess.WorkingDir = workingDir.String sess.ParentID = parentID.String @@ -965,6 +968,7 @@ func (s *SQLiteSessionStore) UpdateSession(ctx context.Context, session *Session Title: session.Title, CreatedAt: session.CreatedAt, ToolsApproved: session.ToolsApproved, + SafetyPolicy: session.SafetyPolicy, HideToolResults: session.HideToolResults, WorkingDir: session.WorkingDir, SendUserMessage: session.SendUserMessage, @@ -996,14 +1000,15 @@ func (s *SQLiteSessionStore) UpdateSession(ctx context.Context, session *Session // Use INSERT OR REPLACE for upsert behavior - creates if not exists, updates if exists _, err = tx.ExecContext(ctx, `INSERT INTO sessions ( - id, tools_approved, input_tokens, output_tokens, title, cost, send_user_message, + 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 ) - VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) + VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) ON CONFLICT(id) DO UPDATE SET title = excluded.title, tools_approved = excluded.tools_approved, + safety_policy = excluded.safety_policy, input_tokens = excluded.input_tokens, output_tokens = excluded.output_tokens, cost = excluded.cost, @@ -1017,7 +1022,7 @@ func (s *SQLiteSessionStore) UpdateSession(ctx context.Context, session *Session thinking = excluded.thinking, parent_id = excluded.parent_id, instruction_context = excluded.instruction_context`, - snapshot.ID, snapshot.ToolsApproved, snapshot.InputTokens, snapshot.OutputTokens, + 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) @@ -1165,12 +1170,12 @@ func (s *SQLiteSessionStore) addSessionTx(ctx context.Context, tx *sql.Tx, sessi _, err = tx.ExecContext(ctx, `INSERT INTO sessions ( - id, tools_approved, input_tokens, output_tokens, title, cost, send_user_message, + 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 ) - VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)`, - session.ID, session.ToolsApproved, session.InputTokens, session.OutputTokens, + 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, From f527431113235ca351f43470420146e41fc77b9c Mon Sep 17 00:00:00 2001 From: Trung Nguyen Date: Thu, 16 Jul 2026 16:02:52 +0200 Subject: [PATCH 2/2] fix(session): also persist approve-tool grants synchronously MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The Persist-on-resume switch symmetry: approve-safe / approve-safer / approve-session mutate session state and are now persisted here. approve-tool mutates session.Permissions.Allow with the same lifetime characteristics, and needs the same treatment or its in-turn grant is silently dropped by the next OnRunStart cycle (or a mid-turn process restart). Skip the persist when toolName is empty — the dispatcher's fallback to the pending tool call's name isn't reachable from here, so the mutation happens in-runtime as before but doesn't get the durable write. --- pkg/server/session_manager.go | 19 +++++++++++-------- 1 file changed, 11 insertions(+), 8 deletions(-) diff --git a/pkg/server/session_manager.go b/pkg/server/session_manager.go index d5a3369ef..a33d817cb 100644 --- a/pkg/server/session_manager.go +++ b/pkg/server/session_manager.go @@ -919,11 +919,8 @@ func (sm *SessionManager) ResumeSession(ctx context.Context, sessionID, confirma return errors.New("session not found") } - // approve-safe / approve-safer / approve-session mutate the session - // mid-turn from the dispatcher, but PersistenceObserver only persists - // on OnRunStart — the mutation would otherwise not survive to the - // next turn. Apply and persist synchronously here so the state is - // durable regardless of when the dispatcher processes the resume. + // Mirror + persist mid-turn session mutations synchronously — + // PersistenceObserver only persists on OnRunStart. if rt.session != nil { mutated := false switch runtime.ResumeType(confirmation) { @@ -936,6 +933,13 @@ func (sm *SessionManager) ResumeSession(ctx context.Context, sessionID, confirma case runtime.ResumeTypeApproveSession: rt.session.SetToolsApproved(true) mutated = true + case runtime.ResumeTypeApproveTool: + // Skip when toolName is empty — the dispatcher's own + // fallback (pending tool call name) isn't reachable here. + if toolName != "" { + rt.session.AppendPermissionAllow(toolName) + mutated = true + } } if mutated { if err := sm.sessionStore.UpdateSession(ctx, rt.session); err != nil { @@ -1156,9 +1160,8 @@ func (sm *SessionManager) SetSessionSafetyPolicy(ctx context.Context, sessionID sm.mux.Lock() defer sm.mux.Unlock() - // Mirror onto the live runtime session (if any) so the running - // dispatcher picks up the new policy on the very next tool call, - // then persist so it survives to the next turn. + // Mirror onto the live runtime session so the dispatcher picks up + // the new policy on the next tool call, not just the next turn. if rt, ok := sm.runtimeSessions.Load(sessionID); ok && rt.session != nil { rt.session.SetSafetyPolicy(policy) return sm.sessionStore.UpdateSession(ctx, rt.session)