diff --git a/pkg/server/session_manager.go b/pkg/server/session_manager.go index 7e0752bd7..a33d817cb 100644 --- a/pkg/server/session_manager.go +++ b/pkg/server/session_manager.go @@ -914,12 +914,41 @@ 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") } + // Mirror + persist mid-turn session mutations synchronously — + // PersistenceObserver only persists on OnRunStart. + 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 + 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 { + 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 +1159,14 @@ func (sm *SessionManager) SetSessionSafetyPolicy(ctx context.Context, sessionID } sm.mux.Lock() defer sm.mux.Unlock() + + // 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) + } + 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,