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
// Chat log — the rendered output stream, exposed via 9P.
chatMu sync.RWMutex
chatLog []byte
chatVers uint32
chatCond *sync.Cond
plan []byte
chatMu sync.RWMutex
chatLog []byte
chatVers uint32
chatCond *sync.Cond
chatSignalMu sync.Mutex
chatSignalCh chan struct{}
plan []byte
}
// Backend returns the active backend from the runtime.
@ -136,6 +138,10 @@ func (ag *Agent) AppendChat(data []byte) {
ag.chatVers++
ag.chatMu.Unlock()
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.
@ -154,9 +160,41 @@ 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
}
// ChatLog returns the raw chat log bytes (caller must hold ChatMu.RLock).
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.
func (ag *Agent) Plan() []byte {
ag.chatMu.RLock()
@ -834,6 +872,7 @@ func NewAgent(cfg AgentParams) *Agent {
state: "idle",
}
ag.signalCh = make(chan struct{})
ag.chatSignalCh = make(chan struct{})
ag.chatCond = sync.NewCond(ag.chatMu.RLocker())
ag.initChatHandler()
ag.SetSessionEnv(cfg.SessionID)

View File

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

View File

@ -14,7 +14,7 @@ import (
func TestNodeOptions(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
}))
if spec.Stream == nil {
@ -265,7 +265,7 @@ func TestOpenLeafBlockOnce(t *testing.T) {
func TestOpenLeafStream(t *testing.T) {
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
})),
)
@ -484,7 +484,7 @@ func TestValidateMultipleBlockingModes(t *testing.T) {
spec := DirNode("/",
FileNode("bad", 0444,
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)