diff --git a/cmd/olliesrv/internal/agent/agent.go b/cmd/olliesrv/internal/agent/agent.go index 5f58622..e79a36e 100644 --- a/cmd/olliesrv/internal/agent/agent.go +++ b/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, diff --git a/cmd/olliesrv/internal/agent/chat.go b/cmd/olliesrv/internal/agent/chat.go index f9b8516..aa5618d 100644 --- a/cmd/olliesrv/internal/agent/chat.go +++ b/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(): diff --git a/cmd/olliesrv/internal/agent/chat_test.go b/cmd/olliesrv/internal/agent/chat_test.go index 877785d..9ed84e3 100644 --- a/cmd/olliesrv/internal/agent/chat_test.go +++ b/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 diff --git a/virtfs/decl.go b/virtfs/decl.go index 677a2a9..389ce14 100644 --- a/virtfs/decl.go +++ b/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 }