Stream assistant text incrementally to log/chat; fix stream lost-wakeup

Two fixes for the rendered text views (chat stream, log snapshot):

1. Lost wakeup: blocking stream readers captured the signal channel
   after their data check, so a write landing in that window signaled a
   now-stale channel and the reader blocked until the next event (the
   TUI's last message appeared only after the next prompt). Stream and
   BlockOnce now capture the signal channel before reading; chat.go
   swaps the channel under chatMu via signalChatLocked and RawLogStream
   captures it while holding the lock.

2. No streaming in TUI: textLog only received finalized blocks, so the
   chat text stream showed nothing until a block completed. SetPartial
   now appends the assistant partial's new suffix to textLog
   incrementally (tracking textPartialLen); AppendBlock appends only the
   trailing separator on finalization to avoid duplicating the body.

cmd/olliesrv/internal/agent/agent.go

@@ -33,8 +33,8 @@ type Agent struct {
 	newBackend    func(string) (backend.Backend, error)
 	currentAction atomic.Pointer[actionHandle]
 	closed        atomic.Bool
-	actionMu      sync.Mutex // serializes action registration with shutdown
-	warnedContext bool       // true after warning about context window limits
+	actionMu      sync.Mutex  // serializes action registration with shutdown
+	warnedContext bool        // true after warning about context window limits
 	resultCache   resultCache // per-agent cache of tool results for read-safe tools
 
 	// Execution state — owned by the agent, protected by stateMu.
@@ -47,10 +47,10 @@ type Agent struct {
 	parentID       string              // immutable ID of the agent that spawned this agent
 	depth          int                 // sub-agent depth (0=top-level, 1=sub-agent, 2=sub-sub-agent)
 	activeChildren atomic.Int32        // number of currently-running child sub-agents
-	peers         map[string]struct{} // peer agent names (same session)
-	peerMu        sync.RWMutex
-	fifo          Fifo // prompt queue
-	toolCallCount atomic.Int64
+	peers          map[string]struct{} // peer agent names (same session)
+	peerMu         sync.RWMutex
+	fifo           Fifo // prompt queue
+	toolCallCount  atomic.Int64
 	pendingInject  atomic.Pointer[string]
 	submitMu       sync.Mutex // serializes Submit calls (commands + turns)
 	stateMu        sync.RWMutex
@@ -62,26 +62,27 @@ type Agent struct {
 	log               *olog.Logger
 	sessionID         string // the owning session's ID
 	memoryWakePending atomic.Bool
-	save              func()                      // trigger debounced persistence
-	flush             func()                      // immediately flush persistence
-	onStateChange     func(agentID, state string) // optional callback for state changes
+	save              func()                                          // trigger debounced persistence
+	flush             func()                                          // immediately flush persistence
+	onStateChange     func(agentID, state string)                     // optional callback for state changes
 	onProcStart       func(agentID string, pid int, tool, cmd string) // optional callback for proc start
-	onClear           func(agentID string)        // optional callback for clear
+	onClear           func(agentID string)                            // optional callback for clear
 
 	// Chat logs — dual format: JSONL for programmatic access, plain text for humans.
-	chatMu       sync.RWMutex
-	rawLog       []byte        // JSONL format (one Block per line)
-	rawStart     int           // byte offset of first byte in rawLog (for truncation accounting)
-	textLog      []byte        // Rendered plain text
-	textStart    int           // byte offset of first byte in textLog
-	chatVers     uint32        // incremented on each append for change detection
-	chatCond     *sync.Cond    // signaled on new chat data
-	chatSignalMu sync.Mutex
-	chatSignalCh chan struct{} // closed on new chat data; replaced with fresh channel
-	partialLine  []byte        // current in-flight (partial) block as a JSONL line; never persisted to rawLog
-	partialVers  uint64        // bumped whenever partialLine changes, for stream change detection
-	plan         []byte        // agent's plan file contents
-	blockCounter uint64        // monotonic counter for block ID generation
+	chatMu         sync.RWMutex
+	rawLog         []byte     // JSONL format (one Block per line)
+	rawStart       int        // byte offset of first byte in rawLog (for truncation accounting)
+	textLog        []byte     // Rendered plain text
+	textStart      int        // byte offset of first byte in textLog
+	chatVers       uint32     // incremented on each append for change detection
+	chatCond       *sync.Cond // signaled on new chat data
+	chatSignalMu   sync.Mutex
+	chatSignalCh   chan struct{} // closed on new chat data; replaced with fresh channel
+	partialLine    []byte        // current in-flight (partial) block as a JSONL line; never persisted to rawLog
+	partialVers    uint64        // bumped whenever partialLine changes, for stream change detection
+	textPartialLen int           // bytes of the current in-flight block's body already written to textLog
+	plan           []byte        // agent's plan file contents
+	blockCounter   uint64        // monotonic counter for block ID generation
 }
 
 // Backend returns the active backend from the runtime.
@@ -323,6 +324,7 @@ func (ag *Agent) Clear() error {
 	ag.textStart = 0
 	ag.partialLine = nil
 	ag.partialVers++
+	ag.textPartialLen = 0
 	ag.chatVers++
 	ag.chatMu.Unlock()
 	ag.chatCond.Broadcast()
@@ -633,14 +635,14 @@ type AgentParams struct {
 // no additional calls are required after construction.
 func NewAgent(cfg AgentParams) *Agent {
 	ag := &Agent{
-		history:         cfg.History,
-		runtime:         cfg.Runtime,
-		profile:         cfg.Profile,
-		name:            cfg.ID, // display name defaults to uname
-		agentsDir:       cfg.AgentsDir,
-		id:              cfg.ID,
-		parentID:        cfg.ParentID,
-		getCwd:          cfg.GetCwd,
+		history:       cfg.History,
+		runtime:       cfg.Runtime,
+		profile:       cfg.Profile,
+		name:          cfg.ID, // display name defaults to uname
+		agentsDir:     cfg.AgentsDir,
+		id:            cfg.ID,
+		parentID:      cfg.ParentID,
+		getCwd:        cfg.GetCwd,
 		systemPrompt:  cfg.SystemPrompt,
 		envBlock:      cfg.EnvBlock,
 		newToolServer: cfg.NewToolServer,

cmd/olliesrv/internal/agent/chat.go

@@ -39,34 +39,55 @@ func (ag *Agent) NextBlockID() string {
 // Partial (in-flight) blocks are never persisted here; see SetPartial.
 func (ag *Agent) AppendBlock(b format.Block) {
 	raw := format.MarshalBlock(b)
-	text := []byte(format.RenderBlock(b))
 
 	ag.chatMu.Lock()
 	ag.rawLog = appendWithOverflow(ag.rawLog, raw, maxLogBytes, &ag.rawStart)
-	ag.textLog = appendWithOverflow(ag.textLog, text, maxLogBytes, &ag.textStart)
+	if ag.textPartialLen > 0 && b.Role == "assistant" && !b.Partial {
+		// Body already streamed to textLog incrementally via SetPartial;
+		// only the trailing separator remains to be appended.
+		ag.textLog = appendWithOverflow(ag.textLog, []byte("\n\n"), maxLogBytes, &ag.textStart)
+	} else {
+		ag.textLog = appendWithOverflow(ag.textLog, []byte(format.RenderBlock(b)), maxLogBytes, &ag.textStart)
+	}
+	ag.textPartialLen = 0
 	// A finalized block supersedes any in-flight partial for the same block.
 	ag.partialLine = nil
 	ag.partialVers++
 	ag.chatVers++
+	ag.signalChatLocked()
 	ag.chatMu.Unlock()
-
-	ag.signalChat()
 }
 
 // SetPartial records the current in-flight (streaming) block as a JSONL line.
 // It is delivered to live stream readers but never appended to rawLog, so the
 // one-shot log.raw snapshot never accumulates per-chunk partial duplicates.
+//
+// For roles that render as visible text (assistant), the newly-arrived suffix
+// is also appended to textLog so the rendered text stream (chat) and snapshot
+// (log) update incrementally instead of only at finalization. textPartialLen
+// tracks how much of the current block's body has already been written to
+// textLog so AppendBlock appends only the trailing separator on finalization.
 func (ag *Agent) SetPartial(b format.Block) {
 	line := format.MarshalBlock(b)
 	ag.chatMu.Lock()
 	ag.partialLine = line
 	ag.partialVers++
+	if b.Role == "assistant" {
+		if len(b.Content) > ag.textPartialLen {
+			delta := b.Content[ag.textPartialLen:]
+			ag.textLog = appendWithOverflow(ag.textLog, []byte(delta), maxLogBytes, &ag.textStart)
+			ag.textPartialLen = len(b.Content)
+		}
+	}
+	ag.signalChatLocked()
 	ag.chatMu.Unlock()
-	ag.signalChat()
 }
 
-// signalChat wakes blocking chat/stream readers.
-func (ag *Agent) signalChat() {
+// signalChatLocked wakes blocking chat/stream readers. The caller must hold
+// chatMu; swapping the signal channel under the same lock that guards the log
+// state closes the lost-wakeup window — a reader that observes no new data
+// under chatMu is guaranteed to be waiting on the channel this call replaces.
+func (ag *Agent) signalChatLocked() {
 	ag.chatCond.Broadcast()
 	ag.chatSignalMu.Lock()
 	close(ag.chatSignalCh)
@@ -152,14 +173,16 @@ func (ag *Agent) RawLogStream(ctx context.Context, base string) ([]byte, string,
 			ag.chatMu.RUnlock()
 			return data, fmt.Sprintf("%d:%d", rawTail, curPv), nil
 		}
-		// 3) Nothing new; sync base and block until signaled.
+		// 3) Nothing new; sync base and block until signaled. Capture the
+		// signal channel while still holding chatMu so a concurrent append
+		// cannot swap it between our data check and our wait (lost wakeup).
 		pv = curPv
 		base = "sync"
-		ag.chatMu.RUnlock()
-
 		ag.chatSignalMu.Lock()
 		sig := ag.chatSignalCh
 		ag.chatSignalMu.Unlock()
+		ag.chatMu.RUnlock()
+
 		select {
 		case <-sig:
 		case <-ctx.Done():

cmd/olliesrv/internal/agent/chat_test.go

@@ -1,8 +1,11 @@
 package agent
 
 import (
+	"context"
+	"strings"
 	"sync"
 	"testing"
+	"time"
 
 	"ollie/format"
 )
@@ -259,6 +262,73 @@ func TestTextLog_RendersHumanReadable(t *testing.T) {
 	}
 }
 
+// TestTextLog_IncrementalPartialNoDuplication verifies that assistant partials
+// stream into textLog incrementally and that finalization does not duplicate the
+// already-streamed body (only the trailing separator is appended).
+func TestTextLog_IncrementalPartialNoDuplication(t *testing.T) {
+	ag := testAgent("session1", "agent1")
+	id := ag.NextBlockID()
+
+	// Stream partials: "Hel" -> "Hello" -> "Hello world"
+	ag.SetPartial(format.Block{Role: "assistant", ID: id, Content: "Hel", Partial: true})
+	if got := string(ag.TextLog()); got != "Hel" {
+		t.Fatalf("after first partial, textLog = %q, want %q", got, "Hel")
+	}
+	ag.SetPartial(format.Block{Role: "assistant", ID: id, Content: "Hello", Partial: true})
+	if got := string(ag.TextLog()); got != "Hello" {
+		t.Fatalf("after second partial, textLog = %q, want %q", got, "Hello")
+	}
+	ag.SetPartial(format.Block{Role: "assistant", ID: id, Content: "Hello world", Partial: true})
+
+	// Finalize.
+	ag.AppendBlock(format.Block{Role: "assistant", ID: id, Content: "Hello world"})
+
+	got := string(ag.TextLog())
+	if got != "Hello world\n\n" {
+		t.Errorf("textLog = %q, want %q (body must not be duplicated)", got, "Hello world\n\n")
+	}
+	if strings.Count(got, "Hello world") != 1 {
+		t.Errorf("body appears %d times, want 1: %q", strings.Count(got, "Hello world"), got)
+	}
+}
+
+// TestRawLogStream_NoLostWakeup verifies that a finalized block appended while
+// a stream reader is between its data check and its wait is still delivered,
+// without requiring a second event. This is the regression test for the TUI /
+// last-message-lost race.
+func TestRawLogStream_NoLostWakeup(t *testing.T) {
+	ag := testAgent("session1", "agent1")
+	ag.AppendBlock(format.Block{Role: "user", ID: "a", Content: "hi"})
+
+	// Catch up: fresh open returns the existing block and a sync base.
+	data, base, err := ag.RawLogStream(context.Background(), "")
+	if err != nil {
+		t.Fatalf("RawLogStream error: %v", err)
+	}
+	if len(data) == 0 {
+		t.Fatal("expected initial history")
+	}
+
+	// Reader is now caught up. Append concurrently while it blocks; the
+	// stream must return the new block promptly (not hang).
+	go func() {
+		ag.AppendBlock(format.Block{Role: "assistant", ID: "b", Content: "last message"})
+	}()
+
+	ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
+	defer cancel()
+	data, _, err = ag.RawLogStream(ctx, base)
+	if err != nil {
+		t.Fatalf("RawLogStream error: %v", err)
+	}
+	if ctx.Err() != nil {
+		t.Fatal("stream blocked: final block was lost (lost wakeup)")
+	}
+	if !contains(string(data), "last message") {
+		t.Errorf("expected final block delivered, got %q", data)
+	}
+}
+
 // Helper to split JSONL
 func splitLines(data []byte) [][]byte {
 	var lines [][]byte

virtfs/decl.go

@@ -130,6 +130,11 @@ func Stream(readFn func(base string) ([]byte, string, error), signal func() <-ch
 	return func(d *FsNodeDecl) {
 		d.Stream = func(ctx context.Context, base string) ([]byte, string, error) {
 			for {
+				// Capture the signal channel BEFORE reading. A writer that
+				// appends between our read and our wait closes this channel,
+				// so the subsequent select returns immediately instead of
+				// blocking on an already-superseded channel (lost wakeup).
+				sig := signal()
 				data, nextBase, err := readFn(base)
 				if err != nil {
 					return nil, base, err
@@ -139,7 +144,7 @@ func Stream(readFn func(base string) ([]byte, string, error), signal func() <-ch
 				}
 				base = nextBase
 				select {
-				case <-signal():
+				case <-sig:
 				case <-ctx.Done():
 					return nil, base, nil
 				}
@@ -161,6 +166,9 @@ func BlockOnce(readFn func() ([]byte, string, error), signal func() <-chan struc
 	return func(d *FsNodeDecl) {
 		d.BlockOnce = func(ctx context.Context, base string) ([]byte, string, error) {
 			for {
+				// Capture the signal channel before reading, closing the
+				// lost-wakeup window (see Stream above).
+				sig := signal()
 				data, hash, err := readFn()
 				if err != nil {
 					return nil, base, err
@@ -172,7 +180,7 @@ func BlockOnce(readFn func() ([]byte, string, error), signal func() <-chan struc
 					return data, hash, nil
 				}
 				select {
-				case <-signal():
+				case <-sig:
 				case <-ctx.Done():
 					return nil, "", nil // timeout: empty response, client re-opens
 				}