virtfs: refactor Stream to accept (readFn, signalFn)

Same pattern as BlockOnce: framework handles the blocking loop.
readFn takes a base string, returns (data, nextBase, error).
signal fires when new data may be available.

Chat stream handlers now use Stream(a.ChatRead, a.ChatSignal).
StreamRaw retained for custom handlers.

streamChat() in support.go is now dead code (replaced by ChatRead).
This commit is contained in:
Levi Neely 2026-08-13 20:50:27 +02:00
parent 794673a5d0
commit dcd119c853
4 changed files with 85 additions and 28 deletions

View File

@ -59,11 +59,13 @@ type Agent struct {
onStateChange func(agentID, state string) // optional callback for state changes onStateChange func(agentID, state string) // optional callback for state changes
// Chat log — the rendered output stream, exposed via 9P. // Chat log — the rendered output stream, exposed via 9P.
chatMu sync.RWMutex chatMu sync.RWMutex
chatLog []byte chatLog []byte
chatVers uint32 chatVers uint32
chatCond *sync.Cond chatCond *sync.Cond
plan []byte chatSignalMu sync.Mutex
chatSignalCh chan struct{}
plan []byte
} }
// Backend returns the active backend from the runtime. // Backend returns the active backend from the runtime.
@ -136,6 +138,10 @@ func (ag *Agent) AppendChat(data []byte) {
ag.chatVers++ ag.chatVers++
ag.chatMu.Unlock() ag.chatMu.Unlock()
ag.chatCond.Broadcast() ag.chatCond.Broadcast()
ag.chatSignalMu.Lock()
close(ag.chatSignalCh)
ag.chatSignalCh = make(chan struct{})
ag.chatSignalMu.Unlock()
} }
// EnsureTrailingNewline appends a newline if the log doesn't end with one. // EnsureTrailingNewline appends a newline if the log doesn't end with one.
@ -154,9 +160,41 @@ func (ag *Agent) ChatMu() *sync.RWMutex { return &ag.chatMu }
// ChatCond returns the condvar for blocking chat readers. // ChatCond returns the condvar for blocking chat readers.
func (ag *Agent) ChatCond() *sync.Cond { return ag.chatCond } 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
}
// ChatLog returns the raw chat log bytes (caller must hold ChatMu.RLock). // ChatLog returns the raw chat log bytes (caller must hold ChatMu.RLock).
func (ag *Agent) ChatLog() []byte { return ag.chatLog } func (ag *Agent) ChatLog() []byte { return ag.chatLog }
// ChatRead returns new chat data since the given offset (base).
// Returns (data, nextBase, error). If no new data, returns empty data.
func (ag *Agent) ChatRead(base string) ([]byte, string, error) {
var offset int
if base != "" {
fmt.Sscanf(base, "%d", &offset)
} else {
ag.chatMu.RLock()
offset = len(ag.chatLog)
ag.chatMu.RUnlock()
}
ag.chatMu.RLock()
log := ag.chatLog
if len(log) <= offset {
ag.chatMu.RUnlock()
return nil, fmt.Sprintf("%d", offset), nil
}
data := make([]byte, len(log)-offset)
copy(data, log[offset:])
newOffset := len(log)
ag.chatMu.RUnlock()
return data, fmt.Sprintf("%d", newOffset), nil
}
// Plan returns a copy of the plan. // Plan returns a copy of the plan.
func (ag *Agent) Plan() []byte { func (ag *Agent) Plan() []byte {
ag.chatMu.RLock() ag.chatMu.RLock()
@ -834,6 +872,7 @@ func NewAgent(cfg AgentParams) *Agent {
state: "idle", state: "idle",
} }
ag.signalCh = make(chan struct{}) ag.signalCh = make(chan struct{})
ag.chatSignalCh = make(chan struct{})
ag.chatCond = sync.NewCond(ag.chatMu.RLocker()) ag.chatCond = sync.NewCond(ag.chatMu.RLocker())
ag.initChatHandler() ag.initChatHandler()
ag.SetSessionEnv(cfg.SessionID) ag.SetSessionEnv(cfg.SessionID)

View File

@ -418,21 +418,17 @@ func buildAgentChildren(a *agent.Agent, s *session.Session) []virtfs.FsNodeDecl
defer a.ChatMu().RUnlock() defer a.ChatMu().RUnlock()
return stripMarkers(a.ChatLog()), nil return stripMarkers(a.ChatLog()), nil
}), }),
virtfs.Stream(func(cctx context.Context, base string) ([]byte, string, error) { virtfs.Stream(func(base string) ([]byte, string, error) {
merged, cancel := mergeCtx(s.Ctx, cctx) data, nextBase, err := a.ChatRead(base)
defer cancel() if err != nil || len(data) == 0 {
for { return data, nextBase, err
data, nextBase, err := streamChat(a, merged, base)
if err != nil || len(data) == 0 {
return data, nextBase, err
}
stripped := stripMarkers(data)
if len(stripped) > 0 {
return stripped, nextBase, nil
}
base = nextBase
} }
}), stripped := stripMarkers(data)
if len(stripped) > 0 {
return stripped, nextBase, nil
}
return nil, nextBase, nil
}, a.ChatSignal),
), ),
virtfs.FileNode("chat.raw", 0444, virtfs.FileNode("chat.raw", 0444,
virtfs.StatOverride(chatStat), virtfs.StatOverride(chatStat),
@ -444,11 +440,7 @@ func buildAgentChildren(a *agent.Agent, s *session.Session) []virtfs.FsNodeDecl
copy(data, log) copy(data, log)
return data, nil return data, nil
}), }),
virtfs.Stream(func(cctx context.Context, base string) ([]byte, string, error) { virtfs.Stream(a.ChatRead, a.ChatSignal),
merged, cancel := mergeCtx(s.Ctx, cctx)
defer cancel()
return streamChat(a, merged, base)
}),
), ),
virtfs.FileNode("statewait", 0444, virtfs.FileNode("statewait", 0444,
virtfs.Read(func() ([]byte, error) { virtfs.Read(func() ([]byte, error) {

View File

@ -123,7 +123,33 @@ func Write(fn func([]byte) error) NodeOption {
} }
// Stream sets a streaming read handler that blocks indefinitely. // Stream sets a streaming read handler that blocks indefinitely.
func Stream(fn func(context.Context, string) ([]byte, string, error)) NodeOption { // readFn returns (data, nextBase, error) given the current base.
// signal returns a channel that is closed when new data may be available.
// The framework blocks on signal until readFn returns non-empty data.
func Stream(readFn func(base string) ([]byte, string, error), signal func() <-chan struct{}) NodeOption {
return func(d *FsNodeDecl) {
d.Stream = func(ctx context.Context, base string) ([]byte, string, error) {
for {
data, nextBase, err := readFn(base)
if err != nil {
return nil, base, err
}
if len(data) > 0 {
return data, nextBase, nil
}
base = nextBase
select {
case <-signal():
case <-ctx.Done():
return nil, base, nil
}
}
}
}
}
// StreamRaw sets a raw streaming read handler with custom blocking logic.
func StreamRaw(fn func(context.Context, string) ([]byte, string, error)) NodeOption {
return func(d *FsNodeDecl) { d.Stream = fn } return func(d *FsNodeDecl) { d.Stream = fn }
} }

View File

@ -14,7 +14,7 @@ import (
func TestNodeOptions(t *testing.T) { func TestNodeOptions(t *testing.T) {
t.Run("Stream", func(t *testing.T) { t.Run("Stream", func(t *testing.T) {
spec := FileNode("test", 0444, Stream(func(ctx context.Context, base string) ([]byte, string, error) { spec := FileNode("test", 0444, StreamRaw(func(ctx context.Context, base string) ([]byte, string, error) {
return nil, "", nil return nil, "", nil
})) }))
if spec.Stream == nil { if spec.Stream == nil {
@ -265,7 +265,7 @@ func TestOpenLeafBlockOnce(t *testing.T) {
func TestOpenLeafStream(t *testing.T) { func TestOpenLeafStream(t *testing.T) {
spec := DirNode("/", spec := DirNode("/",
FileNode("streaming", 0444, Stream(func(ctx context.Context, base string) ([]byte, string, error) { FileNode("streaming", 0444, StreamRaw(func(ctx context.Context, base string) ([]byte, string, error) {
return []byte("stream"), "cursor", nil return []byte("stream"), "cursor", nil
})), })),
) )
@ -484,7 +484,7 @@ func TestValidateMultipleBlockingModes(t *testing.T) {
spec := DirNode("/", spec := DirNode("/",
FileNode("bad", 0444, FileNode("bad", 0444,
BlockOnceRaw(func(context.Context, string) ([]byte, string, error) { return nil, "", nil }), BlockOnceRaw(func(context.Context, string) ([]byte, string, error) { return nil, "", nil }),
Stream(func(context.Context, string) ([]byte, string, error) { return nil, "", nil }), StreamRaw(func(context.Context, string) ([]byte, string, error) { return nil, "", nil }),
), ),
) )
BuildTree(spec) BuildTree(spec)