From aad4f6bce1cace88e862e2c32d9a308eaa499e88 Mon Sep 17 00:00:00 2001 From: Trung Nguyen Date: Tue, 16 Jun 2026 11:18:46 +0200 Subject: [PATCH] feat: add session-scoped plan/build mode Plan mode lets the runtime filter tools to read-only ones and inject a per-turn system reminder, so the agent drafts a plan instead of taking actions. Build mode is the default and preserves today's behaviour. The mode is a per-session property, persisted alongside the rest of the session, and exposed via: - POST /api/sessions { ..., mode } - GET /api/sessions/:id -> { ..., mode } - PATCH /api/sessions/:id/mode { mode } Tool filtering is driven by the MCP-spec ReadOnlyHint annotation, so it extends to user-added MCP tools without any per-tool config. --- pkg/api/types.go | 13 ++++++ pkg/runtime/harness.go | 3 +- pkg/runtime/loop.go | 41 +++++++++++++++++-- pkg/runtime/plan_mode.go | 45 +++++++++++++++++++++ pkg/runtime/runtime_test.go | 54 +++++++++++++++++++++++++ pkg/server/server.go | 22 ++++++++++ pkg/server/server_test.go | 58 +++++++++++++++++++++++++++ pkg/server/session_manager.go | 28 +++++++++++++ pkg/session/branch.go | 1 + pkg/session/migrations.go | 7 ++++ pkg/session/migrations_pinned_test.go | 2 +- pkg/session/session.go | 50 +++++++++++++++++++++++ pkg/session/store.go | 27 ++++++++----- pkg/session/store_test.go | 45 +++++++++++++++++++++ 14 files changed, 379 insertions(+), 17 deletions(-) create mode 100644 pkg/runtime/plan_mode.go diff --git a/pkg/api/types.go b/pkg/api/types.go index 5994ea5fb..ec6aac21b 100644 --- a/pkg/api/types.go +++ b/pkg/api/types.go @@ -135,6 +135,7 @@ type SessionResponse struct { OutputTokens int64 `json:"output_tokens"` WorkingDir string `json:"working_dir,omitempty"` Permissions *session.PermissionsConfig `json:"permissions,omitempty"` + Mode session.Mode `json:"mode,omitempty"` } // UpdateSessionPermissionsRequest represents a request to update session permissions. @@ -142,6 +143,17 @@ type UpdateSessionPermissionsRequest struct { Permissions *session.PermissionsConfig `json:"permissions"` } +// UpdateSessionModeRequest represents a request to update a session's mode. +type UpdateSessionModeRequest struct { + Mode session.Mode `json:"mode"` +} + +// UpdateSessionModeResponse represents the response from updating a session's mode. +type UpdateSessionModeResponse struct { + ID string `json:"id"` + Mode session.Mode `json:"mode"` +} + // ResumeSessionRequest represents a request to resume a session type ResumeSessionRequest struct { Confirmation string `json:"confirmation"` @@ -304,6 +316,7 @@ type SessionSnapshotResponse struct { Messages []session.Message `json:"messages"` ToolsApproved bool `json:"tools_approved"` Permissions *session.PermissionsConfig `json:"permissions,omitempty"` + Mode session.Mode `json:"mode,omitempty"` InputTokens int64 `json:"input_tokens"` OutputTokens int64 `json:"output_tokens"` diff --git a/pkg/runtime/harness.go b/pkg/runtime/harness.go index c48b380f8..ffc0b9c34 100644 --- a/pkg/runtime/harness.go +++ b/pkg/runtime/harness.go @@ -46,7 +46,8 @@ func (r *LocalRuntime) runHarnessAgent(ctx context.Context, sess *session.Sessio }() turnStartMsgs := r.executeTurnStartHooks(ctx, sess, a, events) - messages := sess.GetMessages(a, append(baseExtra, turnStartMsgs...)...) + planReminder := planModeReminderMessages(sess) + messages := sess.GetMessages(a, append(append(baseExtra, turnStartMsgs...), planReminder...)...) stop, msg, rewritten := r.executeBeforeLLMCallHooks(ctx, sess, a, modelID, 1, messages) if stop { slog.WarnContext(ctx, "before_llm_call hook signalled run termination", diff --git a/pkg/runtime/loop.go b/pkg/runtime/loop.go index 59f23b2d1..11696d4b9 100644 --- a/pkg/runtime/loop.go +++ b/pkg/runtime/loop.go @@ -266,7 +266,7 @@ func (r *LocalRuntime) runStreamLoop(ctx context.Context, sess *session.Session, sink.Emit(ErrorWithCode(ErrorCodeToolFailed, fmt.Sprintf("failed to get tools: %v", err))) return } - agentTools = filterExcludedTools(agentTools, sess.ExcludedTools) + agentTools = filterToolsForSession(agentTools, sess) sink.Emit(ToolsetInfo(len(agentTools), false, a.Name())) @@ -348,7 +348,7 @@ func (r *LocalRuntime) runStreamLoop(ctx context.Context, sess *session.Session, sink.Emit(ErrorWithCode(ErrorCodeToolFailed, fmt.Sprintf("failed to get tools: %v", err))) return } - agentTools = filterExcludedTools(agentTools, sess.ExcludedTools) + agentTools = filterToolsForSession(agentTools, sess) // Emit updated tool count. After a ToolListChanged MCP notification // the cache is invalidated, so getTools above re-fetches from the @@ -554,7 +554,13 @@ func (r *LocalRuntime) runTurn( // files) refresh every turn while session-level context (cwd, OS, // arch) stays stable — all without bloating the stored history. turnStartMsgs := r.executeTurnStartHooks(ctx, sess, a, events) - messages := sess.GetMessages(a, slices.Concat(ls.sessionStartMsgs, ls.userPromptMsgs, turnStartMsgs)...) + // Plan-mode reminder rides alongside the turn_start hook output so it + // participates in the same per-turn splice (and the cache_control marker + // that GetMessages applies to the last extra). It is appended last so its + // instruction is the most recent system context the model sees before the + // user prompt — minimising the chance the model ignores it. + planReminder := planModeReminderMessages(sess) + messages := sess.GetMessages(a, slices.Concat(ls.sessionStartMsgs, ls.userPromptMsgs, turnStartMsgs, planReminder)...) slog.DebugContext(ctx, "Retrieved messages for processing", "agent", a.Name(), "message_count", len(messages)) // before_llm_call hooks fire just before the model is invoked. @@ -990,6 +996,33 @@ func filterExcludedTools(agentTools []tools.Tool, excluded []string) []tools.Too return filtered } +// filterToolsForSession applies all session-level tool filters: the explicit +// ExcludedTools name list (used by skill sub-sessions) and, when the session +// is in plan mode, anything whose tool definition doesn't advertise +// ReadOnlyHint. The MCP spec's ReadOnlyHint is the canonical "this tool has +// no side effects" signal, so it's the right knob for plan mode and it +// extends naturally to user-added MCP tools without any per-tool config. +func filterToolsForSession(agentTools []tools.Tool, sess *session.Session) []tools.Tool { + out := filterExcludedTools(agentTools, sess.ExcludedTools) + if sess.Mode == session.ModePlan { + out = filterToReadOnlyTools(out) + } + return out +} + +// filterToReadOnlyTools keeps only tools whose definition advertises +// ReadOnlyHint. Used by plan mode to hide every write/execute tool from the +// model so it can't reach for them even if the system reminder is ignored. +func filterToReadOnlyTools(agentTools []tools.Tool) []tools.Tool { + filtered := make([]tools.Tool, 0, len(agentTools)) + for _, t := range agentTools { + if t.Annotations.ReadOnlyHint { + filtered = append(filtered, t) + } + } + return filtered +} + // reprobe re-runs ensureToolSetsAreStarted after a batch of tool calls. // If new tools became available (by name-set diff), it emits a ToolsetInfo // event to update the TUI immediately. The new tools will be picked up by @@ -1010,7 +1043,7 @@ func (r *LocalRuntime) reprobe( slog.WarnContext(ctx, "reprobe: getTools failed", "agent", a.Name(), "error", err) return } - updated = filterExcludedTools(updated, sess.ExcludedTools) + updated = filterToolsForSession(updated, sess) // Emit any pending warnings that getTools just generated. r.emitAgentWarnings(a, events) diff --git a/pkg/runtime/plan_mode.go b/pkg/runtime/plan_mode.go new file mode 100644 index 000000000..c0a916b77 --- /dev/null +++ b/pkg/runtime/plan_mode.go @@ -0,0 +1,45 @@ +package runtime + +import ( + "github.com/docker/docker-agent/pkg/chat" + "github.com/docker/docker-agent/pkg/session" +) + +// planModeReminder is the per-turn system instruction injected when a session +// is in plan mode. Two layers enforce plan mode: the runtime hides every +// non-read-only tool from the model (see filterToolsForSession in loop.go), +// and this reminder tells the model how it should behave. Hiding the tools +// is the hard guarantee; the reminder is the explanation, so the model +// produces a useful plan instead of just bouncing off missing tools. +const planModeReminder = ` +You are currently in PLAN MODE. + +In this mode you research the codebase, ask clarifying questions, and write a +clear, actionable plan for the user. You MUST NOT make any changes to the +system: + +- No edits to files (no write, edit, create, or delete). +- No shell commands or background jobs. +- No state-changing tool calls of any kind. + +Only read-only tools have been made available to you for this turn. If you try +to call a tool that isn't in your list, the user has explicitly disabled it +for planning. + +End the turn by presenting the plan in your final message and asking the user +to review it. The user will switch you to BUILD MODE when they want execution +to begin. +` + +// planModeReminderMessages returns the system-reminder messages to splice +// before the conversation history when sess is in plan mode. Returns nil for +// other modes so callers can use it unconditionally. +func planModeReminderMessages(sess *session.Session) []chat.Message { + if sess == nil || sess.Mode != session.ModePlan { + return nil + } + return []chat.Message{{ + Role: chat.MessageRoleSystem, + Content: planModeReminder, + }} +} diff --git a/pkg/runtime/runtime_test.go b/pkg/runtime/runtime_test.go index 86cc3fc62..1203ccf83 100644 --- a/pkg/runtime/runtime_test.go +++ b/pkg/runtime/runtime_test.go @@ -2395,6 +2395,60 @@ func TestFilterExcludedTools(t *testing.T) { }) } +func TestFilterToolsForSession_PlanMode(t *testing.T) { + readOnly := tools.Tool{Name: "read_file", Annotations: tools.ToolAnnotations{ReadOnlyHint: true}} + mutating := tools.Tool{Name: "write_file"} + all := []tools.Tool{readOnly, mutating, {Name: "shell"}} + + t.Run("build mode keeps all tools", func(t *testing.T) { + sess := &session.Session{Mode: session.ModeBuild} + result := filterToolsForSession(all, sess) + assert.Len(t, result, 3) + }) + + t.Run("empty mode is treated as build", func(t *testing.T) { + // Sessions loaded before the mode column existed have Mode == "". + sess := &session.Session{} + result := filterToolsForSession(all, sess) + assert.Len(t, result, 3) + }) + + t.Run("plan mode keeps only read-only tools", func(t *testing.T) { + sess := &session.Session{Mode: session.ModePlan} + result := filterToolsForSession(all, sess) + assert.Len(t, result, 1) + assert.Equal(t, "read_file", result[0].Name) + }) + + t.Run("plan mode still respects ExcludedTools", func(t *testing.T) { + readOnly2 := tools.Tool{Name: "list_directory", Annotations: tools.ToolAnnotations{ReadOnlyHint: true}} + sess := &session.Session{ + Mode: session.ModePlan, + ExcludedTools: []string{"read_file"}, + } + result := filterToolsForSession([]tools.Tool{readOnly, readOnly2, mutating}, sess) + assert.Len(t, result, 1) + assert.Equal(t, "list_directory", result[0].Name) + }) +} + +func TestPlanModeReminderMessages(t *testing.T) { + t.Run("build mode returns nil", func(t *testing.T) { + assert.Nil(t, planModeReminderMessages(&session.Session{Mode: session.ModeBuild})) + }) + + t.Run("nil session returns nil", func(t *testing.T) { + assert.Nil(t, planModeReminderMessages(nil)) + }) + + t.Run("plan mode returns a single system reminder", func(t *testing.T) { + msgs := planModeReminderMessages(&session.Session{Mode: session.ModePlan}) + assert.Len(t, msgs, 1) + assert.Equal(t, chat.MessageRoleSystem, msgs[0].Role) + assert.Contains(t, msgs[0].Content, "PLAN MODE") + }) +} + func TestMergeExcludedTools(t *testing.T) { t.Run("both empty", func(t *testing.T) { assert.Nil(t, mergeExcludedTools(nil, nil)) diff --git a/pkg/server/server.go b/pkg/server/server.go index 257c25faf..84475ae14 100644 --- a/pkg/server/server.go +++ b/pkg/server/server.go @@ -71,6 +71,7 @@ func (s *Server) registerRoutes() { group.POST("/sessions/:id/resume", s.resumeSession) group.POST("/sessions/:id/tools/toggle", s.toggleSessionYolo) group.PATCH("/sessions/:id/permissions", s.updateSessionPermissions) + group.PATCH("/sessions/:id/mode", s.updateSessionMode) group.PATCH("/sessions/:id/title", s.updateSessionTitle) group.PATCH("/sessions/:id/tokens", s.updateSessionTokens) group.PATCH("/sessions/:id/starred", s.setSessionStarred) @@ -249,6 +250,7 @@ func (s *Server) getSession(c echo.Context) error { OutputTokens: sess.OutputTokens, WorkingDir: sess.WorkingDir, Permissions: sess.Permissions, + Mode: sess.Mode, }) } @@ -329,6 +331,26 @@ func (s *Server) updateSessionPermissions(c echo.Context) error { return c.JSON(http.StatusOK, map[string]string{"message": "session permissions updated"}) } +func (s *Server) updateSessionMode(c echo.Context) error { + sessionID := c.Param("id") + var req api.UpdateSessionModeRequest + if err := c.Bind(&req); err != nil { + return echo.NewHTTPError(http.StatusBadRequest, fmt.Sprintf("invalid request body: %v", err)) + } + if !req.Mode.IsValid() { + return echo.NewHTTPError(http.StatusBadRequest, fmt.Sprintf("invalid mode %q; must be one of: %s, %s", req.Mode, session.ModeBuild, session.ModePlan)) + } + + if err := s.sm.UpdateSessionMode(c.Request().Context(), sessionID, req.Mode); err != nil { + return echo.NewHTTPError(http.StatusInternalServerError, fmt.Sprintf("failed to update session mode: %v", err)) + } + + return c.JSON(http.StatusOK, api.UpdateSessionModeResponse{ + ID: sessionID, + Mode: req.Mode, + }) +} + func (s *Server) updateSessionTitle(c echo.Context) error { sessionID := c.Param("id") var req api.UpdateSessionTitleRequest diff --git a/pkg/server/server_test.go b/pkg/server/server_test.go index 4bf503544..21146761c 100644 --- a/pkg/server/server_test.go +++ b/pkg/server/server_test.go @@ -199,6 +199,64 @@ func TestServer_UpdateSessionTitle(t *testing.T) { assert.Equal(t, newTitle, sessionResp.Title) } +func TestServer_UpdateSessionMode(t *testing.T) { + t.Parallel() + + ctx := t.Context() + store := session.NewInMemorySessionStore() + lnPath := startServerWithStore(t, ctx, prepareAgentsDir(t), store) + + // Create a session in default (build) mode. + createResp := httpDo(t, ctx, http.MethodPost, lnPath, "/api/sessions", map[string]any{}) + var createdSession session.Session + unmarshal(t, createResp, &createdSession) + require.NotEmpty(t, createdSession.ID) + + // Switch the session into plan mode. + patchResp := httpDo(t, ctx, http.MethodPatch, lnPath, + "/api/sessions/"+createdSession.ID+"/mode", + api.UpdateSessionModeRequest{Mode: session.ModePlan}) + var modeResp api.UpdateSessionModeResponse + unmarshal(t, patchResp, &modeResp) + + assert.Equal(t, createdSession.ID, modeResp.ID) + assert.Equal(t, session.ModePlan, modeResp.Mode) + + // GET should reflect the new mode. + getResp := httpGET(t, ctx, lnPath, "/api/sessions/"+createdSession.ID) + var sessionResp api.SessionResponse + unmarshal(t, getResp, &sessionResp) + assert.Equal(t, session.ModePlan, sessionResp.Mode) + + // Switch back to build mode. + patchResp = httpDo(t, ctx, http.MethodPatch, lnPath, + "/api/sessions/"+createdSession.ID+"/mode", + api.UpdateSessionModeRequest{Mode: session.ModeBuild}) + unmarshal(t, patchResp, &modeResp) + assert.Equal(t, session.ModeBuild, modeResp.Mode) +} + +func TestServer_CreateSession_AcceptsMode(t *testing.T) { + t.Parallel() + + ctx := t.Context() + store := session.NewInMemorySessionStore() + lnPath := startServerWithStore(t, ctx, prepareAgentsDir(t), store) + + // Creating a session with mode=plan should persist that mode. + createResp := httpDo(t, ctx, http.MethodPost, lnPath, "/api/sessions", + map[string]any{"mode": string(session.ModePlan)}) + var createdSession session.Session + unmarshal(t, createResp, &createdSession) + require.NotEmpty(t, createdSession.ID) + assert.Equal(t, session.ModePlan, createdSession.Mode) + + getResp := httpGET(t, ctx, lnPath, "/api/sessions/"+createdSession.ID) + var sessionResp api.SessionResponse + unmarshal(t, getResp, &sessionResp) + assert.Equal(t, session.ModePlan, sessionResp.Mode) +} + func startServerWithStore(t *testing.T, ctx context.Context, agentsDir string, store session.Store) string { t.Helper() diff --git a/pkg/server/session_manager.go b/pkg/server/session_manager.go index ba13671f3..781603a33 100644 --- a/pkg/server/session_manager.go +++ b/pkg/server/session_manager.go @@ -316,6 +316,7 @@ func (sm *SessionManager) GetSessionSnapshot(ctx context.Context, id string) (*a Messages: sess.GetAllMessages(), ToolsApproved: sess.ToolsApproved, Permissions: sess.Permissions, + Mode: sess.Mode, InputTokens: sess.InputTokens, OutputTokens: sess.OutputTokens, Streaming: streaming, @@ -353,6 +354,10 @@ func (sm *SessionManager) CreateSession(ctx context.Context, sessionTemplate *se opts = append(opts, session.WithPermissions(sessionTemplate.Permissions)) } + if sessionTemplate.Mode != "" { + opts = append(opts, session.WithMode(sessionTemplate.Mode)) + } + sess := session.New(opts...) // Copy model-related fields from the template so callers can pin a @@ -741,6 +746,29 @@ func (sm *SessionManager) UpdateSessionPermissions(ctx context.Context, sessionI return sm.sessionStore.UpdateSession(ctx, sess) } +// UpdateSessionMode updates the interaction mode (build/plan) for a session. +// If the session is actively running, it also updates the in-memory session +// object so the next turn's tool filter and plan-mode reminder see the new +// mode without having to round-trip through the store. +func (sm *SessionManager) UpdateSessionMode(ctx context.Context, sessionID string, mode session.Mode) error { + mode = session.NormalizeMode(mode) + sm.mux.Lock() + defer sm.mux.Unlock() + + if rt, ok := sm.runtimeSessions.Load(sessionID); ok && rt.session != nil { + rt.session.Mode = mode + slog.DebugContext(ctx, "Updated mode for active session", "session_id", sessionID, "mode", mode) + return sm.sessionStore.UpdateSession(ctx, rt.session) + } + + sess, err := sm.sessionStore.GetSession(ctx, sessionID) + if err != nil { + return err + } + sess.Mode = mode + return sm.sessionStore.UpdateSession(ctx, sess) +} + // UpdateSessionTitle updates the title for a session. // If the session is actively running, it also updates the in-memory session // object to prevent subsequent runtime saves from overwriting the title. diff --git a/pkg/session/branch.go b/pkg/session/branch.go index 6de7b5ed6..32dc22fa5 100644 --- a/pkg/session/branch.go +++ b/pkg/session/branch.go @@ -76,6 +76,7 @@ func (s *Session) Clone() *Session { CustomModelsUsed: cloneStringSlice(s.CustomModelsUsed), AttachedFiles: cloneStringSlice(s.AttachedFiles), ExcludedTools: cloneStringSlice(s.ExcludedTools), + Mode: s.Mode, AgentName: s.AgentName, ParentID: s.ParentID, MessageUsageHistory: slices.Clone(s.MessageUsageHistory), diff --git a/pkg/session/migrations.go b/pkg/session/migrations.go index 400d41bd1..3cd176964 100644 --- a/pkg/session/migrations.go +++ b/pkg/session/migrations.go @@ -400,6 +400,13 @@ func getAllMigrations() []Migration { Description: "Add first_kept_entry column to session_items for compaction-preserved messages", UpSQL: `ALTER TABLE session_items ADD COLUMN first_kept_entry INTEGER DEFAULT 0`, }, + { + ID: 22, + Name: "022_add_mode_column", + Description: "Add mode column to sessions table for build/plan mode", + UpSQL: `ALTER TABLE sessions ADD COLUMN mode TEXT DEFAULT ''`, + DownSQL: `ALTER TABLE sessions DROP COLUMN mode`, + }, } } diff --git a/pkg/session/migrations_pinned_test.go b/pkg/session/migrations_pinned_test.go index 1ffdff8d0..96ae7c5aa 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 = "0c6d5df46b970104cf49988ee3931e33643d5db85c68dd41b74b639d0094cec9" + const wantDigest = "399611d010efb60b9349257e05e0e68702d432e5019cdd56b4ae4e69654ac691" if got != wantDigest { t.Fatalf(`migration catalogue content has changed. diff --git a/pkg/session/session.go b/pkg/session/session.go index fffdb82e5..90b31a93f 100644 --- a/pkg/session/session.go +++ b/pkg/session/session.go @@ -159,6 +159,13 @@ type Session struct { // recursive run_skill calls. ExcludedTools []string `json:"-"` + // Mode is the session's interaction mode. ModeBuild (default) gives the + // agent its full toolset. ModePlan filters the toolset to read-only tools + // and injects a system reminder so the agent drafts a plan instead of + // making changes. The mode can be flipped at any time via + // PATCH /api/sessions/:id/mode; the next turn picks it up. + Mode Mode `json:"mode,omitempty"` + // AgentName, when set, tells RunStream which agent to use for this session // instead of reading from the shared runtime currentAgent field. This is // required for background agent tasks where multiple sessions may run @@ -185,6 +192,41 @@ type MessageUsageRecord struct { Usage chat.Usage `json:"usage"` } +// Mode is the session's interaction mode (build vs plan). +// +// ModeBuild is the default and gives the agent its full toolset. +// ModePlan filters the toolset to read-only tools (anything whose tool +// definition lacks Annotations.ReadOnlyHint) and injects a per-turn system +// reminder telling the agent to plan rather than act. The runtime applies +// both effects automatically based on this field, so callers only need to +// flip the mode — they don't have to compute tool lists themselves. +type Mode string + +const ( + ModeBuild Mode = "build" + ModePlan Mode = "plan" +) + +// IsValid reports whether m is a known mode. +func (m Mode) IsValid() bool { + switch m { + case ModeBuild, ModePlan: + return true + default: + return false + } +} + +// NormalizeMode returns m if it is a known mode, or ModeBuild otherwise. +// Use this when reading mode from external input (persistence, HTTP body) +// to make sure downstream code always sees a valid mode. +func NormalizeMode(m Mode) Mode { + if m.IsValid() { + return m + } + return ModeBuild +} + // PermissionsConfig defines session-level tool permission overrides // using pattern-based rules (Allow/Ask/Deny arrays). type PermissionsConfig struct { @@ -767,6 +809,14 @@ func WithExcludedTools(names []string) Opt { } } +// WithMode sets the session's interaction mode. An empty or unknown mode is +// normalised to ModeBuild so callers can pass through user input directly. +func WithMode(mode Mode) Opt { + return func(s *Session) { + s.Mode = NormalizeMode(mode) + } +} + // WithAttachedFiles seeds the session with absolute paths of files the user // attached. Used when creating sub-sessions so that delegated agents inherit // the parent's file context. Empty and duplicate paths are dropped. diff --git a/pkg/session/store.go b/pkg/session/store.go index 40970c579..a54ac95db 100644 --- a/pkg/session/store.go +++ b/pkg/session/store.go @@ -230,6 +230,7 @@ func (s *InMemorySessionStore) UpdateSession(_ context.Context, session *Session AgentModelOverrides: cloneStringMap(session.AgentModelOverrides), CustomModelsUsed: cloneStringSlice(session.CustomModelsUsed), AttachedFiles: slices.Clone(session.AttachedFiles), + Mode: session.Mode, ParentID: session.ParentID, } session.mu.RUnlock() @@ -354,7 +355,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` +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, mode` // sessionPersistedFields holds the encoded form of a Session's JSON-bearing // columns plus the SQL representation of parent_id (nil for the empty @@ -595,12 +596,12 @@ func (s *SQLiteSessionStore) AddSession(ctx context.Context, session *Session) e `INSERT INTO sessions ( id, tools_approved, 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 - ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)`, + custom_models_used, thinking, parent_id, mode + ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)`, session.ID, session.ToolsApproved, 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.CustomModelsUsedJSON, false, fields.ParentID, string(session.Mode)) if err != nil { return err } @@ -628,6 +629,7 @@ func scanSession(scanner interface { workingDir sql.NullString permissionsJSON sql.NullString parentID sql.NullString + modeStr sql.NullString agentModelOverridesJSON string customModelsUsedJSON string createdAtStr string @@ -639,6 +641,7 @@ func scanSession(scanner interface { &sess.Title, &sess.Cost, &sess.SendUserMessage, &sess.MaxIterations, &workingDir, &createdAtStr, &sess.Starred, &permissionsJSON, &agentModelOverridesJSON, &customModelsUsedJSON, &thinking, &parentID, + &modeStr, ) if err != nil { return nil, err @@ -651,6 +654,7 @@ func scanSession(scanner interface { sess.WorkingDir = workingDir.String sess.ParentID = parentID.String + sess.Mode = NormalizeMode(Mode(modeStr.String)) if permissionsJSON.Valid && permissionsJSON.String != "" { sess.Permissions = &PermissionsConfig{} @@ -908,9 +912,9 @@ func (s *SQLiteSessionStore) UpdateSession(ctx context.Context, session *Session `INSERT INTO sessions ( 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 + custom_models_used, thinking, parent_id, mode ) - VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) + VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) ON CONFLICT(id) DO UPDATE SET title = excluded.title, tools_approved = excluded.tools_approved, @@ -925,11 +929,12 @@ func (s *SQLiteSessionStore) UpdateSession(ctx context.Context, session *Session agent_model_overrides = excluded.agent_model_overrides, custom_models_used = excluded.custom_models_used, thinking = excluded.thinking, - parent_id = excluded.parent_id`, + parent_id = excluded.parent_id, + mode = excluded.mode`, session.ID, session.ToolsApproved, 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.CustomModelsUsedJSON, false, fields.ParentID, string(session.Mode)) if err != nil { return err } @@ -1076,14 +1081,14 @@ func (s *SQLiteSessionStore) addSessionTx(ctx context.Context, tx *sql.Tx, sessi `INSERT INTO sessions ( 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 + custom_models_used, thinking, parent_id, mode ) - VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)`, + VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)`, session.ID, session.ToolsApproved, 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.ParentID, string(session.Mode)) return err } diff --git a/pkg/session/store_test.go b/pkg/session/store_test.go index cef3d0479..2b66673e2 100644 --- a/pkg/session/store_test.go +++ b/pkg/session/store_test.go @@ -542,6 +542,51 @@ func TestUpdateSession_Permissions(t *testing.T) { assert.Equal(t, []string{"dangerous_*"}, retrieved.Permissions.Deny) } +func TestSessionMode_SQLite(t *testing.T) { + tempDB := filepath.Join(t.TempDir(), "test_session_mode.db") + + store, err := NewSQLiteSessionStore(tempDB) + require.NoError(t, err) + defer store.(*SQLiteSessionStore).Close() + + // Default Mode (empty string) round-trips as ModeBuild after scan. + defaultSess := &Session{ + ID: "default-mode-session", + Title: "Default mode", + CreatedAt: time.Now(), + } + require.NoError(t, store.AddSession(t.Context(), defaultSess)) + retrieved, err := store.GetSession(t.Context(), defaultSess.ID) + require.NoError(t, err) + assert.Equal(t, ModeBuild, retrieved.Mode) + + // ModePlan persists and reloads. + planSess := &Session{ + ID: "plan-mode-session", + Title: "Plan mode", + CreatedAt: time.Now(), + Mode: ModePlan, + } + require.NoError(t, store.AddSession(t.Context(), planSess)) + retrieved, err = store.GetSession(t.Context(), planSess.ID) + require.NoError(t, err) + assert.Equal(t, ModePlan, retrieved.Mode) + + // Mode can be flipped via UpdateSession. + planSess.Mode = ModeBuild + require.NoError(t, store.UpdateSession(t.Context(), planSess)) + retrieved, err = store.GetSession(t.Context(), planSess.ID) + require.NoError(t, err) + assert.Equal(t, ModeBuild, retrieved.Mode) +} + +func TestNormalizeMode(t *testing.T) { + assert.Equal(t, ModeBuild, NormalizeMode("")) + assert.Equal(t, ModeBuild, NormalizeMode("garbage")) + assert.Equal(t, ModeBuild, NormalizeMode(ModeBuild)) + assert.Equal(t, ModePlan, NormalizeMode(ModePlan)) +} + func TestAgentModelOverrides_SQLite(t *testing.T) { tempDB := filepath.Join(t.TempDir(), "test_model_overrides.db")