Skip to content

fix(session): persist SafetyPolicy across turns #3685

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
dgageot merged 2 commits into from
Jul 17, 2026
+61 −13
Merged

fix(session): persist SafetyPolicy across turns #3685

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
458b780
fix(session): persist SafetyPolicy across turns
trungutt Jul 16, 2026
f527431
fix(session): also persist approve-tool grants synchronously
trungutt Jul 16, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
39 changes: 38 additions & 1 deletion pkg/server/session_manager.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
}
}
Comment thread
trungutt marked this conversation as resolved.
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,
Expand Down Expand Up @@ -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
Expand Down
6 changes: 6 additions & 0 deletions pkg/session/migrations.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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 ''`,
},
}
}

Expand Down
2 changes: 1 addition & 1 deletion pkg/session/migrations_pinned_test.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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.

Expand Down
27 changes: 16 additions & 11 deletions pkg/session/store.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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)
Expand Down Expand Up @@ -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
Expand All @@ -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,
Expand All @@ -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

Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -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,
Expand All @@ -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)
Expand Down Expand Up @@ -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,
Expand Down
Loading