Skip to content

Commit bf7ca85

Browse files
dgageotdgageot
authored
Merge pull request #3685 from trungutt/persist-safety-policy
fix(session): persist SafetyPolicy across turns
2 parents a47d1ab + f527431 commit bf7ca85

4 files changed

Lines changed: 61 additions & 13 deletions

File tree

‎pkg/server/session_manager.go‎

Lines changed: 38 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -914,12 +914,41 @@ func (sm *SessionManager) ResumeSession(ctx context.Context, sessionID, confirma
914914
sm.mux.Lock()
915915
defer sm.mux.Unlock()
916916

917-
// Ensure the session runtime exists
918917
rt, exists := sm.runtimeSessions.Load(sessionID)
919918
if !exists {
920919
return errors.New("session not found")
921920
}
922921

922+
// Mirror + persist mid-turn session mutations synchronously —
923+
// PersistenceObserver only persists on OnRunStart.
924+
if rt.session != nil {
925+
mutated := false
926+
switch runtime.ResumeType(confirmation) {
927+
case runtime.ResumeTypeApproveSafe:
928+
rt.session.SetSafetyPolicy(session.SafetyPolicySafeAuto)
929+
mutated = true
930+
case runtime.ResumeTypeApproveSafer:
931+
rt.session.SetSafetyPolicy(session.SafetyPolicySafer)
932+
mutated = true
933+
case runtime.ResumeTypeApproveSession:
934+
rt.session.SetToolsApproved(true)
935+
mutated = true
936+
case runtime.ResumeTypeApproveTool:
937+
// Skip when toolName is empty — the dispatcher's own
938+
// fallback (pending tool call name) isn't reachable here.
939+
if toolName != "" {
940+
rt.session.AppendPermissionAllow(toolName)
941+
mutated = true
942+
}
943+
}
944+
if mutated {
945+
if err := sm.sessionStore.UpdateSession(ctx, rt.session); err != nil {
946+
slog.WarnContext(ctx, "failed to persist mid-turn session state",
947+
"session_id", sessionID, "confirmation", confirmation, "err", err)
948+
}
949+
}
950+
}
951+
923952
rt.runtime.Resume(ctx, runtime.ResumeRequest{
924953
Type: runtime.ResumeType(confirmation),
925954
Reason: reason,
@@ -1130,6 +1159,14 @@ func (sm *SessionManager) SetSessionSafetyPolicy(ctx context.Context, sessionID
11301159
}
11311160
sm.mux.Lock()
11321161
defer sm.mux.Unlock()
1162+
1163+
// Mirror onto the live runtime session so the dispatcher picks up
1164+
// the new policy on the next tool call, not just the next turn.
1165+
if rt, ok := sm.runtimeSessions.Load(sessionID); ok && rt.session != nil {
1166+
rt.session.SetSafetyPolicy(policy)
1167+
return sm.sessionStore.UpdateSession(ctx, rt.session)
1168+
}
1169+
11331170
sess, err := sm.sessionStore.GetSession(ctx, sessionID)
11341171
if err != nil {
11351172
return err

‎pkg/session/migrations.go‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -413,6 +413,12 @@ func getAllMigrations() []Migration {
413413
Description: "Add cost column to session_items so compaction summary costs survive reload",
414414
UpSQL: `ALTER TABLE session_items ADD COLUMN cost REAL NOT NULL DEFAULT 0`,
415415
},
416+
{
417+
ID: 24,
418+
Name: "024_add_safety_policy_column",
419+
Description: "Add safety_policy column to sessions table",
420+
UpSQL: `ALTER TABLE sessions ADD COLUMN safety_policy TEXT DEFAULT ''`,
421+
},
416422
}
417423
}
418424

‎pkg/session/migrations_pinned_test.go‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -39,7 +39,7 @@ func TestMigrationCatalogIsContentPinned(t *testing.T) {
3939

4040
got := digestMigrationCatalog(getAllMigrations())
4141

42-
const wantDigest = "19cc1ae9c44d6ba716afea1f0f7f1a87623caf0f4540a751613e988ee2fd7049"
42+
const wantDigest = "ba6212344ecacb721bda5a00a5f60ad02a891b79d8a12103484252f78b597569"
4343
if got != wantDigest {
4444
t.Fatalf(`migration catalogue content has changed.
4545

‎pkg/session/store.go‎

Lines changed: 16 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -226,6 +226,7 @@ func (s *InMemorySessionStore) UpdateSession(_ context.Context, session *Session
226226
Evals: session.Evals,
227227
CreatedAt: session.CreatedAt,
228228
ToolsApproved: session.ToolsApproved,
229+
SafetyPolicy: session.SafetyPolicy,
229230
HideToolResults: session.HideToolResults,
230231
WorkingDir: session.WorkingDir,
231232
SendUserMessage: session.SendUserMessage,
@@ -376,7 +377,7 @@ type SQLiteSessionStore struct {
376377
// sessionSelectColumns is the canonical SELECT list for the sessions table.
377378
// The column order matches what scanSession expects; all read paths use this
378379
// constant so that adding a column requires updating exactly one place.
379-
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`
380+
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`
380381

381382
// sessionPersistedFields holds the encoded form of a Session's JSON-bearing
382383
// 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
643644

644645
_, err = tx.ExecContext(ctx,
645646
`INSERT INTO sessions (
646-
id, tools_approved, input_tokens, output_tokens, title, cost, send_user_message,
647+
id, tools_approved, safety_policy, input_tokens, output_tokens, title, cost, send_user_message,
647648
max_iterations, working_dir, created_at, permissions, agent_model_overrides,
648649
custom_models_used, thinking, parent_id, instruction_context
649-
) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)`,
650-
session.ID, session.ToolsApproved, session.InputTokens, session.OutputTokens, session.Title,
650+
) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)`,
651+
session.ID, session.ToolsApproved, string(session.SafetyPolicy), session.InputTokens, session.OutputTokens, session.Title,
651652
session.Cost, session.SendUserMessage, session.MaxIterations, session.WorkingDir,
652653
session.CreatedAt.Format(time.RFC3339), fields.PermissionsJSON, fields.AgentModelOverridesJSON,
653654
fields.CustomModelsUsedJSON, false, fields.ParentID, fields.InstructionContextJSON)
@@ -675,6 +676,7 @@ func scanSession(scanner interface {
675676
) (*Session, error) {
676677
var (
677678
sess Session
679+
safetyPolicy sql.NullString
678680
workingDir sql.NullString
679681
permissionsJSON sql.NullString
680682
parentID sql.NullString
@@ -686,7 +688,7 @@ func scanSession(scanner interface {
686688
)
687689

688690
err := scanner.Scan(
689-
&sess.ID, &sess.ToolsApproved, &sess.InputTokens, &sess.OutputTokens,
691+
&sess.ID, &sess.ToolsApproved, &safetyPolicy, &sess.InputTokens, &sess.OutputTokens,
690692
&sess.Title, &sess.Cost, &sess.SendUserMessage, &sess.MaxIterations,
691693
&workingDir, &createdAtStr, &sess.Starred, &permissionsJSON,
692694
&agentModelOverridesJSON, &customModelsUsedJSON, &thinking, &parentID, &instructionContextJSON,
@@ -696,6 +698,7 @@ func scanSession(scanner interface {
696698
}
697699

698700
sess.CreatedAt = parseCreatedAt(createdAtStr)
701+
sess.SafetyPolicy = SafetyPolicy(safetyPolicy.String)
699702
sess.WorkingDir = workingDir.String
700703
sess.ParentID = parentID.String
701704

@@ -965,6 +968,7 @@ func (s *SQLiteSessionStore) UpdateSession(ctx context.Context, session *Session
965968
Title: session.Title,
966969
CreatedAt: session.CreatedAt,
967970
ToolsApproved: session.ToolsApproved,
971+
SafetyPolicy: session.SafetyPolicy,
968972
HideToolResults: session.HideToolResults,
969973
WorkingDir: session.WorkingDir,
970974
SendUserMessage: session.SendUserMessage,
@@ -996,14 +1000,15 @@ func (s *SQLiteSessionStore) UpdateSession(ctx context.Context, session *Session
9961000
// Use INSERT OR REPLACE for upsert behavior - creates if not exists, updates if exists
9971001
_, err = tx.ExecContext(ctx,
9981002
`INSERT INTO sessions (
999-
id, tools_approved, input_tokens, output_tokens, title, cost, send_user_message,
1003+
id, tools_approved, safety_policy, input_tokens, output_tokens, title, cost, send_user_message,
10001004
max_iterations, working_dir, created_at, starred, permissions, agent_model_overrides,
10011005
custom_models_used, thinking, parent_id, instruction_context
10021006
)
1003-
VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
1007+
VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
10041008
ON CONFLICT(id) DO UPDATE SET
10051009
title = excluded.title,
10061010
tools_approved = excluded.tools_approved,
1011+
safety_policy = excluded.safety_policy,
10071012
input_tokens = excluded.input_tokens,
10081013
output_tokens = excluded.output_tokens,
10091014
cost = excluded.cost,
@@ -1017,7 +1022,7 @@ func (s *SQLiteSessionStore) UpdateSession(ctx context.Context, session *Session
10171022
thinking = excluded.thinking,
10181023
parent_id = excluded.parent_id,
10191024
instruction_context = excluded.instruction_context`,
1020-
snapshot.ID, snapshot.ToolsApproved, snapshot.InputTokens, snapshot.OutputTokens,
1025+
snapshot.ID, snapshot.ToolsApproved, string(snapshot.SafetyPolicy), snapshot.InputTokens, snapshot.OutputTokens,
10211026
snapshot.Title, snapshot.Cost, snapshot.SendUserMessage, snapshot.MaxIterations, snapshot.WorkingDir,
10221027
snapshot.CreatedAt.Format(time.RFC3339), snapshot.Starred, fields.PermissionsJSON, fields.AgentModelOverridesJSON,
10231028
fields.CustomModelsUsedJSON, false, fields.ParentID, fields.InstructionContextJSON)
@@ -1165,12 +1170,12 @@ func (s *SQLiteSessionStore) addSessionTx(ctx context.Context, tx *sql.Tx, sessi
11651170

11661171
_, err = tx.ExecContext(ctx,
11671172
`INSERT INTO sessions (
1168-
id, tools_approved, input_tokens, output_tokens, title, cost, send_user_message,
1173+
id, tools_approved, safety_policy, input_tokens, output_tokens, title, cost, send_user_message,
11691174
max_iterations, working_dir, created_at, starred, permissions, agent_model_overrides,
11701175
custom_models_used, thinking, parent_id, instruction_context
11711176
)
1172-
VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)`,
1173-
session.ID, session.ToolsApproved, session.InputTokens, session.OutputTokens,
1177+
VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)`,
1178+
session.ID, session.ToolsApproved, string(session.SafetyPolicy), session.InputTokens, session.OutputTokens,
11741179
session.Title, session.Cost, session.SendUserMessage, session.MaxIterations,
11751180
session.WorkingDir, session.CreatedAt.Format(time.RFC3339), session.Starred,
11761181
fields.PermissionsJSON, fields.AgentModelOverridesJSON, fields.CustomModelsUsedJSON, false,

0 commit comments

Comments
 (0)