test message

Author: Patel230

cmd/iterate/config.go

@@ -16,6 +16,7 @@ import (
 type iterConfig struct {
 	Provider      string   `json:"provider"       toml:"provider"`
 	Model         string   `json:"model"          toml:"model"`
+	APIKey        string   `json:"api_key,omitempty"        toml:"api_key"`
 	OllamaBaseURL string   `json:"ollama_base_url,omitempty" toml:"ollama_base_url"`
 	SafeMode      bool     `json:"safe_mode,omitempty"      toml:"safe_mode"`
 	DeniedTools   []string `json:"denied_tools,omitempty"   toml:"denied_tools"`
@@ -34,6 +35,19 @@ type iterConfig struct {
 	DenyDirs  []string `json:"deny_dirs,omitempty"  toml:"deny_dirs"`
 }
 
+func (c *iterConfig) GetAPIKey() string         { return c.APIKey }
+func (c *iterConfig) SetAPIKey(v string)        { c.APIKey = v }
+func (c *iterConfig) SetProvider(v string)      { c.Provider = v }
+func (c *iterConfig) SetModel(v string)         { c.Model = v }
+func (c *iterConfig) SetNotify(v bool)          { c.Notify = v }
+func (c *iterConfig) SetSafeMode(v bool)        { c.SafeMode = v }
+func (c *iterConfig) SetTheme(v string)         { c.Theme = v }
+func (c *iterConfig) SetThinkingLevel(v string) { c.ThinkingLevel = v }
+func (c *iterConfig) SetTemperature(v float64)  { c.Temperature = v }
+func (c *iterConfig) SetMaxTokens(v int)        { c.MaxTokens = v }
+func (c *iterConfig) SetCacheEnabled(v bool)    { c.CacheEnabled = v }
+func (c *iterConfig) SetOllamaBaseURL(v string) { c.OllamaBaseURL = v }
+
 func configPath() string {
 	home, _ := os.UserHomeDir()
 	return filepath.Join(home, ".iterate", "config.json")
@@ -130,7 +144,7 @@ func saveConfig(cfg iterConfig) {
 		}
 		var buf bytes.Buffer
 		if err := toml.NewEncoder(&buf).Encode(cfg); err == nil {
-			if err := os.WriteFile(tomlPath, buf.Bytes(), 0o644); err != nil {
+			if err := atomicWriteFile(tomlPath, buf.Bytes(), 0o644); err != nil {
 				slog.Warn("failed to write TOML config", "err", err)
 			}
 			return
@@ -141,7 +155,7 @@ func saveConfig(cfg iterConfig) {
 		slog.Warn("failed to create JSON config dir", "err", err)
 	}
 	data, _ := json.MarshalIndent(cfg, "", "  ")
-	if err := os.WriteFile(jsonPath, data, 0o644); err != nil {
+	if err := atomicWriteFile(jsonPath, data, 0o644); err != nil {
 		slog.Warn("failed to write JSON config", "err", err)
 	}
 }

cmd/iterate/features_sessions.go

@@ -60,7 +60,7 @@ func saveSession(name string, messages []iteragent.Message) error {
 	if err != nil {
 		return err
 	}
-	return os.WriteFile(path, data, 0o644)
+	return atomicWriteFile(path, data, 0o644)
 }
 
 func loadSession(name string) ([]iteragent.Message, error) {
@@ -117,7 +117,7 @@ func loadBookmarks() []Bookmark {
 
 func saveBookmarks(bms []Bookmark) {
 	data, _ := json.MarshalIndent(bms, "", "  ")
-	if err := os.WriteFile(bookmarksPath(), data, 0o644); err != nil {
+	if err := atomicWriteFile(bookmarksPath(), data, 0o644); err != nil {
 		slog.Warn("failed to write bookmarks file", "err", err)
 	}
 }

cmd/iterate/fileutil.go

@@ -0,0 +1,39 @@
+package main
+
+import (
+	"fmt"
+	"os"
+	"path/filepath"
+)
+
+// atomicWriteFile writes data to path atomically by writing to a temp file in
+// the same directory then renaming it. This prevents partial writes from
+// corrupting the file if the process is interrupted mid-write.
+func atomicWriteFile(path string, data []byte, perm os.FileMode) error {
+	dir := filepath.Dir(path)
+	tmp, err := os.CreateTemp(dir, ".tmp-")
+	if err != nil {
+		return fmt.Errorf("create temp file: %w", err)
+	}
+	tmpName := tmp.Name()
+	defer func() {
+		// Clean up temp file if rename didn't happen.
+		_ = os.Remove(tmpName)
+	}()
+
+	if _, err := tmp.Write(data); err != nil {
+		_ = tmp.Close()
+		return fmt.Errorf("write temp file: %w", err)
+	}
+	if err := tmp.Chmod(perm); err != nil {
+		_ = tmp.Close()
+		return fmt.Errorf("chmod temp file: %w", err)
+	}
+	if err := tmp.Close(); err != nil {
+		return fmt.Errorf("close temp file: %w", err)
+	}
+	if err := os.Rename(tmpName, path); err != nil {
+		return fmt.Errorf("rename temp file to %s: %w", path, err)
+	}
+	return nil
+}

cmd/iterate/main_mode.go

@@ -16,9 +16,10 @@ import (
 
 func runMode(ctx context.Context, f mainFlags, absRepo string, logger *slog.Logger) {
 	cfg := loadConfig()
-	providerName, modelName := resolveProviderConfig(f.provider, f.model, cfg)
+	providerName, modelName, apiKey := resolveProviderConfig(f.provider, f.model, f.apiKey, cfg)
 	f.provider = providerName
 	f.model = modelName
+	f.apiKey = apiKey
 
 	p, err := initProvider(f.provider, f.apiKey, logger)
 	if err != nil {

cmd/iterate/memory_project.go

@@ -51,7 +51,7 @@ func saveProjectMemory(repoPath string, m projectMemory) error {
 	if err != nil {
 		return err
 	}
-	return os.WriteFile(path, data, 0o644)
+	return atomicWriteFile(path, data, 0o644)
 }
 
 func addProjectMemoryNote(repoPath, note string) error {

cmd/iterate/provider.go

@@ -9,24 +9,28 @@ import (
 
 // resolveProviderConfig merges flag values with persisted config.
 // Flags take precedence: only defaults ("gemini", empty) are overridden.
-func resolveProviderConfig(flagProvider, flagModel string, cfg iterConfig) (provider, model string) {
+func resolveProviderConfig(flagProvider, flagModel, flagAPIKey string, cfg iterConfig) (provider, model, apiKey string) {
 	provider = flagProvider
 	model = flagModel
+	apiKey = flagAPIKey
 
 	if provider == "gemini" && cfg.Provider != "" && cfg.Provider != "gemini" {
 		provider = cfg.Provider
 	}
 	if model == "" && cfg.Model != "" {
 		model = cfg.Model
 	}
+	if apiKey == "" && cfg.APIKey != "" {
+		apiKey = cfg.APIKey
+	}
 	if cfg.OllamaBaseURL != "" && os.Getenv("OLLAMA_BASE_URL") == "" {
 		os.Setenv("OLLAMA_BASE_URL", cfg.OllamaBaseURL)
 	}
 	if model != "" {
 		os.Setenv("ITERATE_MODEL", model)
 	}
 
-	return provider, model
+	return provider, model, apiKey
 }
 
 // resolveThinkingLevel returns the effective thinking level, falling back to

cmd/iterate/repl.go

@@ -41,6 +41,14 @@ var replRegistry = commands.DefaultRegistry()
 // iterateVersion is the current version string.
 const iterateVersion = "dev"
 
+// closeProvider shuts down a provider if it implements io.Closer.
+func closeProvider(p iteragent.Provider) {
+	type closer interface{ Close() error }
+	if c, ok := p.(closer); ok {
+		_ = c.Close()
+	}
+}
+
 func makeAgent(p iteragent.Provider, repoPath string, thinking iteragent.ThinkingLevel, logger *slog.Logger) *iteragent.Agent {
 	base := iteragent.DefaultTools(repoPath)
 	switch currentMode {
@@ -181,7 +189,10 @@ func runREPL(ctx context.Context, p iteragent.Provider, repoPath string, thinkin
 	thinking = initREPL(repoPath, thinking)
 
 	a := makeAgent(p, repoPath, thinking, logger)
-	defer func() { _ = a.Close() }() // best-effort cleanup
+	defer func() {
+		_ = a.Close()
+		closeProvider(p)
+	}()
 
 	printHeader(p, thinking, repoPath)
 
@@ -262,6 +273,7 @@ func handleModelProviderSwitch(line string, p *iteragent.Provider, thinking *ite
 				slog.Error("provider switch failed", "provider", providerName, "error", err)
 				fmt.Printf("%serror: %s%s\n\n", colorRed, err, colorReset)
 			} else {
+				closeProvider(*p)
 				*p = newP
 				os.Setenv("ITERATE_PROVIDER", providerName)
 				_ = (*a).Close() // best-effort cleanup
@@ -422,6 +434,17 @@ func buildCommandContext(repoPath, line string, parts []string, p iteragent.Prov
 		InputHistory:        selector.InputHistoryRef,
 		StopWatch:           stopWatch,
 		Pool:                agentPool,
+		Config: commands.ConfigCallbacks{
+			LoadConfig: func() interface{} {
+				c := loadConfig()
+				return &c
+			},
+			SaveConfig: func(v interface{}) {
+				if c, ok := v.(*iterConfig); ok {
+					saveConfig(*c)
+				}
+			},
+		},
 		Session: commands.SessionCallbacks{
 			SaveSession:   saveSession,
 			LoadSession:   loadSession,

cmd/iterate/repl_streaming.go

@@ -102,6 +102,9 @@ func logTokenDelta(beforeTokens int) {
 func streamAndPrint(ctx context.Context, a *iteragent.Agent, prompt string, repoPath string) {
 	recordMessage()
 
+	// Sync pinned messages into the agent before each request.
+	a.SetPinnedMessages(getPinnedMessages())
+
 	reqCtx, cancel := context.WithCancel(ctx)
 	sess.RequestCancel = cancel
 	defer func() {
@@ -118,9 +121,13 @@ func streamAndPrint(ctx context.Context, a *iteragent.Agent, prompt string, repo
 
 	var fullContent string
 	var toolStart time.Time
+	var ttft time.Duration
 	beforeTokens := sess.Tokens
 
 	for e := range events {
+		if ttft == 0 && iteragent.EventType(e.Type) == iteragent.EventTokenUpdate && e.Content != "" {
+			ttft = time.Since(start).Round(time.Millisecond)
+		}
 		fullContent, toolStart = processStreamEvent(e, fullContent, toolStart, stopOnce, newSpinner, repoPath)
 	}
 	a.Finish()
@@ -135,7 +142,12 @@ func streamAndPrint(ctx context.Context, a *iteragent.Agent, prompt string, repo
 	elapsed := time.Since(start).Round(time.Millisecond)
 
 	updateSessionTokens(a, fullContent)
-	printFinalStats(elapsed, beforeTokens, fullContent)
+	printFinalStats(elapsed, ttft, beforeTokens, fullContent)
+
+	// Autosave after each turn so a crash doesn't lose the session.
+	if len(a.Messages) > 0 {
+		_ = saveSession("autosave", a.Messages)
+	}
 }
 
 // newSpinnerController creates a spinner control pair (stopOnce, newSpinner).
@@ -147,6 +159,7 @@ func newSpinnerController() (func(), func(string)) {
 		stopOnce     func()
 		spinnerLabel string
 	)
+	stopOnce = func() {} // no-op until a spinner is started
 	newSpinner := func(label string) {
 		spinnerLabel = label
 		stopSpinner = make(chan struct{})
@@ -161,7 +174,9 @@ func newSpinnerController() (func(), func(string)) {
 		go spinner(stopSpinner, spinnerDone, spinnerLabel)
 	}
 	_ = spinnerLabel
-	return stopOnce, newSpinner
+	// Return a wrapper that always calls the *current* stopOnce, not the
+	// initial no-op captured at return time.
+	return func() { stopOnce() }, newSpinner
 }
 
 // processStreamEvent handles a single agent stream event and returns updated state.
@@ -184,7 +199,15 @@ func processStreamEvent(e iteragent.Event, fullContent string, toolStart time.Ti
 		fmt.Printf("%s%s %s%s", col, icon, label, colorReset)
 	case iteragent.EventToolExecutionEnd:
 		elapsed := time.Since(toolStart).Round(time.Millisecond)
-		printToolResult(e.Result, elapsed, e.ToolName, repoPath)
+		if e.IsError {
+			fmt.Printf("%s  ✗ %s%s  %s%s%s\n",
+				colorDim, colorReset,
+				colorRed, e.Result, colorReset,
+				colorDim)
+			fmt.Print(colorReset)
+		} else {
+			printToolResult(e.Result, elapsed, e.ToolName, repoPath)
+		}
 		newSpinner("thinking")
 	case iteragent.EventContextCompacted:
 		fmt.Printf("\r\033[K%s[context compacted]%s\n", colorDim, colorReset)
@@ -202,34 +225,38 @@ func printToolResult(result string, elapsed time.Duration, toolName string, repo
 	}
 }
 
-// updateSessionTokens updates session token counters from the last agent message.
+// updateSessionTokens updates session token counters from the last assistant
+// message with usage data. Searches backwards since tool result messages (role
+// user) may appear after the final assistant message.
 func updateSessionTokens(a *iteragent.Agent, fullContent string) {
-	if len(a.Messages) > 0 {
-		last := a.Messages[len(a.Messages)-1]
-		if last.Usage != nil {
-			sess.InputTokens += last.Usage.InputTokens
-			sess.OutputTokens += last.Usage.OutputTokens
-			sess.CacheRead += last.Usage.CacheRead
-			sess.CacheWrite += last.Usage.CacheWrite
-			sess.Tokens += last.Usage.TotalTokens
+	for i := len(a.Messages) - 1; i >= 0; i-- {
+		if a.Messages[i].Usage != nil {
+			u := a.Messages[i].Usage
+			sess.InputTokens += u.InputTokens
+			sess.OutputTokens += u.OutputTokens
+			sess.CacheRead += u.CacheRead
+			sess.CacheWrite += u.CacheWrite
+			sess.Tokens += u.TotalTokens
+			return
 		}
-	} else {
-		approxTokens := len(fullContent) / 4
-		sess.Tokens += approxTokens
-		sess.OutputTokens += approxTokens
 	}
+	// Fallback: approximate from streamed content length.
+	approxTokens := len(fullContent) / 4
+	sess.Tokens += approxTokens
+	sess.OutputTokens += approxTokens
 }
 
-// printFinalStats prints token delta, status line, and debug log.
-func printFinalStats(elapsed time.Duration, beforeTokens int, fullContent string) {
-	fmt.Println()
-	logTokenDelta(beforeTokens)
+// printFinalStats prints the status line and debug log.
+func printFinalStats(elapsed, ttft time.Duration, beforeTokens int, fullContent string) {
+	delta := sess.Tokens - beforeTokens
+
 	fmt.Println()
 	selector.InputTokens = sess.InputTokens
 	selector.OutputTokens = sess.OutputTokens
 	selector.SafeMode = cfg.SafeMode
-	selector.PrintStatusLine(elapsed)
+	selector.TTFT = ttft
+	selector.PrintStatusLine(elapsed, delta)
 	fmt.Println()
 
-	slog.Debug("request completed", "elapsed_ms", elapsed.Milliseconds(), "response_chars", len(fullContent), "total_tokens", sess.Tokens)
+	slog.Debug("request completed", "elapsed_ms", elapsed.Milliseconds(), "ttft_ms", ttft.Milliseconds(), "response_chars", len(fullContent), "total_tokens", sess.Tokens)
 }

go.mod

@@ -11,6 +11,8 @@ require (
 	golang.org/x/tools v0.43.0
 )
 
+replace github.com/GrayCodeAI/iteragent => ../iteragent
+
 require (
 	github.com/google/go-querystring v1.2.0 // indirect
 	golang.org/x/mod v0.34.0 // indirect

go.sum

@@ -1,9 +1,5 @@
 github.com/BurntSushi/toml v1.6.0 h1:dRaEfpa2VI55EwlIW72hMRHdWouJeRF7TPYhI+AUQjk=
 github.com/BurntSushi/toml v1.6.0/go.mod h1:ukJfTF/6rtPPRCnwkur4qwRxa8vTRFBF0uk2lLoLwho=
-github.com/GrayCodeAI/iteragent v1.4.0 h1:BSW/Zj0142sXO6vZrC47LUqLt5O3M6eepgpG+VEDMxk=
-github.com/GrayCodeAI/iteragent v1.4.0/go.mod h1:wsIv7o17FthxfJ1zzbqgPdnVSSlUOxmEC2ch96S+4hY=
-github.com/GrayCodeAI/iteragent v1.5.0 h1:Z7weM5U6SOgFBf6AH2clSxPGeXB0URxGXF1kfhqhpG4=
-github.com/GrayCodeAI/iteragent v1.5.0/go.mod h1:wsIv7o17FthxfJ1zzbqgPdnVSSlUOxmEC2ch96S+4hY=
 github.com/google/go-cmp v0.6.0 h1:ofyhxvXcZhMsU5ulbFiLKl/XBFqE1GSq7atu8tAmTRI=
 github.com/google/go-cmp v0.6.0/go.mod h1:17dUlkBOakJ0+DkrSSNjCkIjxS6bF9zb3elmeNGIjoY=
 github.com/google/go-github/v61 v61.0.0 h1:VwQCBwhyE9JclCI+22/7mLB1PuU9eowCXKY5pNlu1go=