ollie/cmd/olliesrv/internal/agent/chat.go

// chat.go — Chat log storage and streaming.
//
// Two logs are maintained:
// - rawLog: JSONL format (one Block per line) for programmatic access
// - textLog: Rendered plain text for human reading
//
// Both are append-only buffers with overflow handling.

package agent

import (
	"bytes"
	"context"
	"crypto/sha256"
	"encoding/hex"
	"fmt"
	"sync"
	"sync/atomic"

	"ollie/format"
)

const maxLogBytes = 16 * 1024 * 1024

// NextBlockID returns the next deterministic block ID and increments the counter.
// Block IDs are sha256(sessionID + agentID + counter)[:8].
func (ag *Agent) NextBlockID() string {
	n := atomic.AddUint64(&ag.blockCounter, 1) - 1
	h := sha256.New()
	h.Write([]byte(ag.sessionID))
	h.Write([]byte(":"))
	h.Write([]byte(ag.id))
	h.Write([]byte(":"))
	h.Write([]byte(fmt.Sprintf("%d", n)))
	return hex.EncodeToString(h.Sum(nil))[:8]
}

// AppendBlock writes a finalized block to both logs and notifies stream readers.
// Partial (in-flight) blocks are never persisted here; see SetPartial.
func (ag *Agent) AppendBlock(b format.Block) {
	raw := format.MarshalBlock(b)

	ag.chatMu.Lock()
	ag.rawLog = appendWithOverflow(ag.rawLog, raw, maxLogBytes, &ag.rawStart)
	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()
}

// 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()
}

// 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)
	ag.chatSignalCh = make(chan struct{})
	ag.chatSignalMu.Unlock()
}

func appendWithOverflow(log, data []byte, max int, start *int) []byte {
	log = append(log, data...)
	if len(log) > max {
		drop := len(log) - max
		log = append([]byte(nil), log[drop:]...)
		*start += drop
	}
	return log
}

// ChatMu returns the chat log mutex for external locking (streaming).
func (ag *Agent) ChatMu() *sync.RWMutex { return &ag.chatMu }

// ChatCond returns the condvar for blocking chat readers.
func (ag *Agent) ChatCond() *sync.Cond { return ag.chatCond }

// ChatSignal returns the current chat signal channel (closed on new chat data).
func (ag *Agent) ChatSignal() <-chan struct{} {
	ag.chatSignalMu.Lock()
	ch := ag.chatSignalCh
	ag.chatSignalMu.Unlock()
	return ch
}

// RawLog returns the JSONL log bytes (caller must hold ChatMu.RLock).
func (ag *Agent) RawLog() []byte { return ag.rawLog }

// TextLog returns the rendered text log bytes (caller must hold ChatMu.RLock).
func (ag *Agent) TextLog() []byte { return ag.textLog }

// RawLogRead returns new JSONL data since the given offset.
func (ag *Agent) RawLogRead(base string) ([]byte, string, error) {
	return ag.logRead(base, &ag.rawLog, &ag.rawStart)
}

// RawLogStream is the authoritative live JSONL source for GUI clients.
//
// On a fresh open (base == "") it replays the finalized history (rawLog, one
// line per block, no partials), then streams live deltas: the current in-flight
// partial block whenever it changes, and finalized blocks as they are appended.
// rawLog is append-only and never rewritten, so finals are delivered with
// offset semantics; partials are delivered out-of-band and collapse by block ID
// on the client. This stream never replays accumulated partials, so reconnects
// stay O(history) rather than O(partial chunks).
//
// base encodes "<rawOffset>:<partialVers>".
func (ag *Agent) RawLogStream(ctx context.Context, base string) ([]byte, string, error) {
	var off int
	var pv uint64
	if base != "" {
		fmt.Sscanf(base, "%d:%d", &off, &pv)
	}
	for {
		ag.chatMu.RLock()
		if base == "" {
			off = ag.rawStart
		}
		if off < ag.rawStart {
			off = ag.rawStart
		}
		rawTail := ag.rawStart + len(ag.rawLog)
		// 1) Deliver any un-sent finalized blocks first (offset-safe).
		if off < rawTail {
			localOff := off - ag.rawStart
			data := make([]byte, len(ag.rawLog)-localOff)
			copy(data, ag.rawLog[localOff:])
			nextPv := ag.partialVers
			ag.chatMu.RUnlock()
			return data, fmt.Sprintf("%d:%d", rawTail, nextPv), nil
		}
		// 2) Caught up to finals: deliver the current partial if it changed.
		curPv := ag.partialVers
		if curPv != pv && len(ag.partialLine) > 0 {
			data := make([]byte, len(ag.partialLine))
			copy(data, ag.partialLine)
			ag.chatMu.RUnlock()
			return data, fmt.Sprintf("%d:%d", rawTail, curPv), nil
		}
		// 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.chatSignalMu.Lock()
		sig := ag.chatSignalCh
		ag.chatSignalMu.Unlock()
		ag.chatMu.RUnlock()

		select {
		case <-sig:
		case <-ctx.Done():
			return nil, fmt.Sprintf("%d:%d", off, pv), nil
		}
	}
}

// TextLogRead returns new rendered text since the given offset.
func (ag *Agent) TextLogRead(base string) ([]byte, string, error) {
	return ag.logRead(base, &ag.textLog, &ag.textStart)
}

func (ag *Agent) logRead(base string, log *[]byte, start *int) ([]byte, string, error) {
	var offset int
	if base != "" {
		fmt.Sscanf(base, "%d", &offset)
	} else {
		// Empty base = start from beginning (for replay)
		ag.chatMu.RLock()
		offset = *start
		ag.chatMu.RUnlock()
	}
	ag.chatMu.RLock()
	if offset < *start {
		offset = *start
	}
	localOffset := offset - *start
	if len(*log) <= localOffset {
		ag.chatMu.RUnlock()
		return nil, fmt.Sprintf("%d", offset), nil
	}
	data := make([]byte, len(*log)-localOffset)
	copy(data, (*log)[localOffset:])
	newOffset := *start + len(*log)
	ag.chatMu.RUnlock()
	return data, fmt.Sprintf("%d", newOffset), nil
}

// BlockByID returns the block with the given ID from the raw log.
func (ag *Agent) BlockByID(blockID string) (format.Block, bool) {
	if blockID == "" {
		return format.Block{}, false
	}
	ag.chatMu.RLock()
	raw := ag.rawLog
	ag.chatMu.RUnlock()

	// Scan JSONL lines for matching ID
	for _, line := range bytes.Split(raw, []byte("\n")) {
		if len(line) == 0 {
			continue
		}
		b, err := format.UnmarshalBlock(line)
		if err != nil {
			continue
		}
		if b.ID == blockID && !b.Partial {
			return b, true
		}
	}
	return format.Block{}, false
}

// --- Plan methods ---

// Plan returns a copy of the plan.
func (ag *Agent) Plan() []byte {
	ag.chatMu.RLock()
	p := make([]byte, len(ag.plan))
	copy(p, ag.plan)
	ag.chatMu.RUnlock()
	return p
}

// SetPlan replaces the plan.
func (ag *Agent) SetPlan(data []byte) {
	ag.chatMu.Lock()
	ag.plan = make([]byte, len(data))
	copy(ag.plan, data)
	ag.chatMu.Unlock()
}