refactor: dedupe cross-package logic and remove dead code from audit

cmd/root.go

@@ -12,7 +12,6 @@ import (
 
 	tea "charm.land/bubbletea/v2"
 	"github.com/mark3labs/kit/internal/app"
-	"github.com/mark3labs/kit/internal/auth"
 	"github.com/mark3labs/kit/internal/config"
 	"github.com/mark3labs/kit/internal/extensions"
 	"github.com/mark3labs/kit/internal/models"
@@ -677,8 +676,8 @@ func globalShortcutsProviderForUI(k *kit.Kit) func() map[string]func() {
 	}
 }
 
-func runNormalMode(ctx context.Context) error {
-	// Validate flag combinations
+// validateModeFlags rejects invalid flag combinations for the root command.
+func validateModeFlags() error {
 	if quietFlag && positionalPrompt == "" {
 		return fmt.Errorf("--quiet requires a prompt (e.g. kit \"your question\" --quiet)")
 	}
@@ -691,21 +690,14 @@ func runNormalMode(ctx context.Context) error {
 	if noExitFlag && positionalPrompt == "" {
 		return fmt.Errorf("--no-exit requires a prompt (e.g. kit \"your question\" --no-exit)")
 	}
+	return nil
+}
 
-	// Set up logging
-	if debugMode {
-		log.SetFlags(log.LstdFlags | log.Lshortfile)
-	}
-
-	// Update debug mode from viper
-	if viper.GetBool("debug") && !debugMode {
-		debugMode = viper.GetBool("debug")
-		log.SetFlags(log.LstdFlags | log.Lshortfile)
-	}
-
-	// Restore persisted model preference when no explicit --model flag or
-	// config file model is set. Precedence: CLI flag > config file > saved
-	// preference > built-in default. This mirrors how themes are persisted.
+// restorePersistedPreferences applies saved model / thinking-level
+// preferences into viper when neither a CLI flag nor a config-file value
+// takes precedence. Precedence: CLI flag > config file > saved preference >
+// built-in default. This mirrors how themes are persisted.
+func restorePersistedPreferences() {
 	// Skip custom/* models unless --provider-url is also provided, since the
 	// custom provider requires a URL that was only valid for the previous session.
 	if !modelFlagChanged && !viper.InConfig("model") {
@@ -724,14 +716,23 @@ func runNormalMode(ctx context.Context) error {
 			viper.Set("thinking-level", pref)
 		}
 	}
+}
+
+// applyProviderURLRouting rewrites the model in viper when --provider-url
+// is set, routing requests through the "custom" (OpenAI-compatible)
+// provider. Must run after restorePersistedPreferences.
+func applyProviderURLRouting() {
+	if viper.GetString("provider-url") == "" {
+		return
+	}
 
 	// When --provider-url is set but no explicit --model was provided,
 	// default to "custom/custom" so the user doesn't need to remember a
 	// provider/model pair for custom OpenAI-compatible endpoints.
 	// This intentionally overrides saved preferences but respects config-file
 	// models — if you specify a model in ~/.kit.yml, it will be used with
 	// custom/custom's provider routing.
-	if viper.GetString("provider-url") != "" && !modelFlagChanged && !viper.InConfig("model") {
+	if !modelFlagChanged && !viper.InConfig("model") {
 		viper.Set("model", "custom/custom")
 	}
 
@@ -746,7 +747,7 @@ func runNormalMode(ctx context.Context) error {
 	// to point a non-OpenAI wire (Anthropic, Google, ...) at a proxy URL,
 	// use the explicit `custom/<name>` form to opt out of the rewrite by
 	// configuring the proxy as that provider in your config file instead.
-	if viper.GetString("provider-url") != "" && modelFlagChanged {
+	if modelFlagChanged {
 		model := viper.GetString("model")
 		if model != "" {
 			name := model
@@ -758,6 +759,26 @@ func runNormalMode(ctx context.Context) error {
 			}
 		}
 	}
+}
+
+func runNormalMode(ctx context.Context) error {
+	if err := validateModeFlags(); err != nil {
+		return err
+	}
+
+	// Set up logging
+	if debugMode {
+		log.SetFlags(log.LstdFlags | log.Lshortfile)
+	}
+
+	// Update debug mode from viper
+	if viper.GetBool("debug") && !debugMode {
+		debugMode = viper.GetBool("debug")
+		log.SetFlags(log.LstdFlags | log.Lshortfile)
+	}
+
+	restorePersistedPreferences()
+	applyProviderURLRouting()
 
 	// Load MCP configuration.
 	mcpConfig, err := config.LoadAndValidateConfig()
@@ -1164,23 +1185,7 @@ func runNormalMode(ctx context.Context) error {
 		// NotifyModelChanged calls prog.Send() which deadlocks. The UI layer
 		// updates m.providerName and m.modelName directly after setModel returns.
 		// Update usage tracker with new model info for correct token counting.
-		if usageTracker != nil {
-			newProvider, newModel, _ := models.ParseModelString(modelString)
-			if newProvider != "unknown" && newModel != "unknown" && newProvider != "ollama" {
-				registry := models.GetGlobalRegistry()
-				if modelInfo := registry.LookupModel(newProvider, newModel); modelInfo != nil {
-					// Check OAuth status for Anthropic models
-					isOAuth := false
-					if newProvider == "anthropic" {
-						_, source, err := auth.GetAnthropicAPIKey(viper.GetString("provider-api-key"))
-						if err == nil && strings.HasPrefix(source, "stored OAuth") {
-							isOAuth = true
-						}
-					}
-					usageTracker.UpdateModelInfo(modelInfo, newProvider, isOAuth)
-				}
-			}
-		}
+		ui.UpdateUsageTrackerForModel(usageTracker, modelString, viper.GetString("provider-api-key"))
 		return nil
 	}
 	emitModelChangeForUI := func(newModel, previousModel, source string) {

internal/acpserver/session.go

@@ -73,111 +73,70 @@ func (r *sessionRegistry) create(ctx context.Context, cwd string) (*acpSession,
 
 	// Wire extension context with headless implementations so extensions
 	// work in ACP mode. TUI-dependent features (widgets, prompts, editor)
-	// become no-ops or return cancelled; all data/model/tool APIs work
-	// identically to interactive mode.
+	// become no-ops or return cancelled; all data/model/tool APIs come from
+	// extbridge.BaseContext and work identically to interactive mode.
 	if kitInstance.Extensions().HasExtensions() {
-		kitInstance.Extensions().SetContext(extensions.Context{
-			SessionID:   sessionID,
-			CWD:         cwd,
-			Model:       kitInstance.GetModelString(),
-			Interactive: false,
-
-			// Output — route through structured logger.
-			Print:      func(text string) { log.Debug("extension: print", "text", text) },
-			PrintInfo:  func(text string) { log.Info("extension: info", "text", text) },
-			PrintError: func(text string) { log.Error("extension: error", "text", text) },
-			PrintBlock: func(opts extensions.PrintBlockOpts) {
-				log.Info("extension: block", "subtitle", opts.Subtitle, "text", opts.Text)
-			},
-
-			// Message injection — no-ops for now; ACP clients drive prompts.
-			SendMessage:   func(string) {},
-			CancelAndSend: func(string) {},
-			Exit:          func() {},
-
-			// TUI widgets/chrome — silent no-ops (no TUI in ACP).
-			SetWidget:       func(extensions.WidgetConfig) {},
-			RemoveWidget:    func(string) {},
-			SetHeader:       func(extensions.HeaderFooterConfig) {},
-			RemoveHeader:    func() {},
-			SetFooter:       func(extensions.HeaderFooterConfig) {},
-			RemoveFooter:    func() {},
-			SetEditor:       func(extensions.EditorConfig) {},
-			ResetEditor:     func() {},
-			SetEditorText:   func(string) {},
-			SetUIVisibility: func(extensions.UIVisibility) {},
-			SetStatus:       func(string, string, int) {},
-			RemoveStatus:    func(string) {},
-
-			// Interactive prompts — return cancelled (no user to prompt).
-			PromptSelect: func(extensions.PromptSelectConfig) extensions.PromptSelectResult {
-				return extensions.PromptSelectResult{Cancelled: true}
-			},
-			PromptConfirm: func(extensions.PromptConfirmConfig) extensions.PromptConfirmResult {
-				return extensions.PromptConfirmResult{Cancelled: true}
-			},
-			PromptInput: func(extensions.PromptInputConfig) extensions.PromptInputResult {
-				return extensions.PromptInputResult{Cancelled: true}
-			},
-			ShowOverlay: func(extensions.OverlayConfig) extensions.OverlayResult {
-				return extensions.OverlayResult{Cancelled: true, Index: -1}
-			},
-			SuspendTUI: func(callback func()) error { callback(); return nil },
-
-			// Data access — delegate to Kit instance.
-			GetContextStats: func() extensions.ContextStats {
-				s := kitInstance.GetContextStats()
-				return extensions.ContextStats{
-					EstimatedTokens: s.EstimatedTokens,
-					ContextLimit:    s.ContextLimit,
-					UsagePercent:    s.UsagePercent,
-					MessageCount:    s.MessageCount,
-				}
-			},
-			GetMessages:    func() []extensions.SessionMessage { return kitInstance.Extensions().GetSessionMessages() },
-			GetSessionPath: func() string { return kitInstance.GetSessionPath() },
-			AppendEntry: func(entryType, data string) (string, error) {
-				return kitInstance.Extensions().AppendEntry(entryType, data)
-			},
-			GetEntries: func(entryType string) []extensions.ExtensionEntry {
-				return kitInstance.Extensions().GetEntries(entryType)
-			},
-
-			// Options, model, and tool management.
-			GetOption: func(name string) string { return kitInstance.Extensions().GetOption(name) },
-			SetOption: func(name, value string) { kitInstance.Extensions().SetOption(name, value) },
-			SetModel: func(modelString string) error {
-				previousModel := kitInstance.Extensions().GetContext().Model
-				if err := kitInstance.SetModel(context.Background(), modelString); err != nil {
-					return err
-				}
-				kitInstance.Extensions().UpdateContextModel(modelString)
-				kitInstance.Extensions().EmitModelChange(modelString, previousModel, "extension")
-				return nil
-			},
-			GetAvailableModels: func() []extensions.ModelInfoEntry { return kitInstance.GetAvailableModels() },
-			EmitCustomEvent:    func(name, data string) { kitInstance.Extensions().EmitCustomEvent(name, data) },
-			GetAllTools:        func() []extensions.ToolInfo { return kitInstance.Extensions().GetToolInfos() },
-			SetActiveTools:     func(names []string) { kitInstance.Extensions().SetActiveTools(names) },
-
-			// LLM completions and subagents.
-			Complete: func(req extensions.CompleteRequest) (extensions.CompleteResponse, error) {
-				return kitInstance.ExecuteCompletion(context.Background(), req)
-			},
-			SpawnSubagent: func(config extensions.SubagentConfig) (*extensions.SubagentHandle, *extensions.SubagentResult, error) {
-				return extbridge.SpawnSubagent(context.Background(), kitInstance, config)
-			},
-
-			// Render — fall back to logging.
-			RenderMessage: func(name, content string) {
-				renderer := kitInstance.Extensions().GetMessageRenderer(name)
-				if renderer != nil && renderer.Render != nil {
-					content = renderer.Render(content, 80)
-				}
-				log.Info("extension: message", "renderer", name, "content", content)
-			},
-			ReloadExtensions: func() error { return kitInstance.Extensions().Reload() },
-		})
+		// Use a background context for subagent spawns: the create() ctx is
+		// request-scoped and may be cancelled before extensions spawn anything.
+		ec := extbridge.BaseContext(context.Background(), kitInstance)
+
+		ec.SessionID = sessionID
+		ec.CWD = cwd
+		ec.Model = kitInstance.GetModelString()
+		ec.Interactive = false
+
+		// Output — route through structured logger.
+		ec.Print = func(text string) { log.Debug("extension: print", "text", text) }
+		ec.PrintInfo = func(text string) { log.Info("extension: info", "text", text) }
+		ec.PrintError = func(text string) { log.Error("extension: error", "text", text) }
+		ec.PrintBlock = func(opts extensions.PrintBlockOpts) {
+			log.Info("extension: block", "subtitle", opts.Subtitle, "text", opts.Text)
+		}
+
+		// Message injection — no-ops for now; ACP clients drive prompts.
+		ec.SendMessage = func(string) {}
+		ec.CancelAndSend = func(string) {}
+		ec.Exit = func() {}
+
+		// TUI widgets/chrome — silent no-ops (no TUI in ACP).
+		ec.SetWidget = func(extensions.WidgetConfig) {}
+		ec.RemoveWidget = func(string) {}
+		ec.SetHeader = func(extensions.HeaderFooterConfig) {}
+		ec.RemoveHeader = func() {}
+		ec.SetFooter = func(extensions.HeaderFooterConfig) {}
+		ec.RemoveFooter = func() {}
+		ec.SetEditor = func(extensions.EditorConfig) {}
+		ec.ResetEditor = func() {}
+		ec.SetEditorText = func(string) {}
+		ec.SetUIVisibility = func(extensions.UIVisibility) {}
+		ec.SetStatus = func(string, string, int) {}
+		ec.RemoveStatus = func(string) {}
+
+		// Interactive prompts — return cancelled (no user to prompt).
+		ec.PromptSelect = func(extensions.PromptSelectConfig) extensions.PromptSelectResult {
+			return extensions.PromptSelectResult{Cancelled: true}
+		}
+		ec.PromptConfirm = func(extensions.PromptConfirmConfig) extensions.PromptConfirmResult {
+			return extensions.PromptConfirmResult{Cancelled: true}
+		}
+		ec.PromptInput = func(extensions.PromptInputConfig) extensions.PromptInputResult {
+			return extensions.PromptInputResult{Cancelled: true}
+		}
+		ec.ShowOverlay = func(extensions.OverlayConfig) extensions.OverlayResult {
+			return extensions.OverlayResult{Cancelled: true, Index: -1}
+		}
+		ec.SuspendTUI = func(callback func()) error { callback(); return nil }
+
+		// Render — fall back to logging.
+		ec.RenderMessage = func(name, content string) {
+			renderer := kitInstance.Extensions().GetMessageRenderer(name)
+			if renderer != nil && renderer.Render != nil {
+				content = renderer.Render(content, 80)
+			}
+			log.Info("extension: message", "renderer", name, "content", content)
+		}
+
+		kitInstance.Extensions().SetContext(ec)
 		kitInstance.Extensions().EmitSessionStart()
 	}
 

internal/app/messages.go

@@ -13,11 +13,6 @@ type MessageStore struct {
 	messages []kit.LLMMessage
 }
 
-// NewMessageStore creates an empty MessageStore.
-func NewMessageStore() *MessageStore {
-	return &MessageStore{}
-}
-
 // NewMessageStoreWithMessages creates a MessageStore pre-populated with the
 // given messages. This is used when loading an existing session at startup.
 func NewMessageStoreWithMessages(msgs []kit.LLMMessage) *MessageStore {