diff --git a/store/session.go b/store/session.go index 279484d..710c3a3 100644 --- a/store/session.go +++ b/store/session.go @@ -26,13 +26,17 @@ type Session struct { cancel context.CancelFunc log []byte logVers uint32 + logMu sync.Mutex + logCond *sync.Cond ChatOffset int plan []byte prevPrompt []byte // last submitted prompt; overwritten on each new submission } func NewSession(id string, core agent.Core, ctx context.Context, cancel context.CancelFunc) *Session { - return &Session{id: id, Core: core, Ctx: ctx, cancel: cancel} + s := &Session{id: id, Core: core, Ctx: ctx, cancel: cancel} + s.logCond = sync.NewCond(&s.logMu) + return s } func (sess *Session) RunnableID() string { return sess.id } @@ -54,16 +58,59 @@ func (sess *Session) AppendLog(data []byte) { sess.log = append(sess.log, data...) sess.logVers++ sess.mu.Unlock() + sess.logCond.Broadcast() } // EnsureTrailingNewline appends a newline if the log doesn't already end with one. +// Always signals waiters to unblock chatwait at end of turn. func (sess *Session) EnsureTrailingNewline() { sess.mu.Lock() if len(sess.log) > 0 && sess.log[len(sess.log)-1] != '\n' { sess.log = append(sess.log, '\n') - sess.logVers++ } + sess.logVers++ sess.mu.Unlock() + sess.logCond.Broadcast() +} + +// WaitLog blocks until the log grows past offset, then returns the new bytes. +// Returns nil if ctx is cancelled or the log is signalled without growth +// (e.g. end of turn). +func (sess *Session) WaitLog(ctx context.Context, offset int) []byte { + stop := context.AfterFunc(ctx, func() { + sess.logCond.Broadcast() + }) + defer stop() + + sess.logMu.Lock() + defer sess.logMu.Unlock() + + // Snapshot version before waiting. + sess.mu.RLock() + startVers := sess.logVers + sess.mu.RUnlock() + + for ctx.Err() == nil { + sess.mu.RLock() + n := len(sess.log) + v := sess.logVers + var data []byte + if n > offset { + data = make([]byte, n-offset) + copy(data, sess.log[offset:]) + } + sess.mu.RUnlock() + if data != nil { + return data + } + // If version changed but no new content (e.g. end-of-turn signal), + // return nil to unblock the caller. + if v != startVers { + return nil + } + sess.logCond.Wait() + } + return nil } // LogInfo returns the current log length and version atomically. diff --git a/store/sessionfile.go b/store/sessionfile.go index fd86c46..58ed225 100644 --- a/store/sessionfile.go +++ b/store/sessionfile.go @@ -35,6 +35,7 @@ var SessionFileList = []struct { {"env", 0444}, {"tail", 0555}, {"prompt.prev", 0444}, + {"chatwait", 0444}, {"context", 0444}, } @@ -192,6 +193,25 @@ func (h *sessionHelper) fileSpec(name string, mode os.FileMode) FileSpec { return []byte(v + "\n"), nil } } + if name == "chatwait" { + fs.Wait = func(connCtx context.Context, base string) ([]byte, error) { + ctx, cancel := context.WithCancel(connCtx) + defer cancel() + context.AfterFunc(h.sess.Ctx, cancel) + offset := 0 + if base == "" { + n, _ := h.sess.LogInfo() + offset = n + } else { + offset, _ = strconv.Atoi(base) + } + data := h.sess.WaitLog(ctx, offset) + if data == nil { + return nil, nil + } + return data, nil + } + } return fs } diff --git a/store/store_test.go b/store/store_test.go index 7fa7fa4..0a800ff 100644 --- a/store/store_test.go +++ b/store/store_test.go @@ -448,6 +448,60 @@ func TestSessionAppendLog(t *testing.T) { } } +func TestSessionWaitLog(t *testing.T) { + sess := testSession("wl") + defer sess.Cancel() + + sess.AppendLog([]byte("existing")) + + // WaitLog should return new content appended after offset. + done := make(chan []byte, 1) + go func() { + data := sess.WaitLog(context.Background(), 8) // offset = len("existing") + done <- data + }() + + // Give goroutine time to block. + time.Sleep(10 * time.Millisecond) + sess.AppendLog([]byte(" new")) + + select { + case data := <-done: + if string(data) != " new" { + t.Errorf("WaitLog = %q; want %q", data, " new") + } + case <-time.After(time.Second): + t.Fatal("WaitLog did not unblock") + } +} + +func TestSessionWaitLogEndOfTurn(t *testing.T) { + sess := testSession("wl2") + defer sess.Cancel() + + sess.AppendLog([]byte("content\n")) + + // WaitLog should unblock on EnsureTrailingNewline (end-of-turn signal) + // even when no new content is added. + done := make(chan []byte, 1) + go func() { + data := sess.WaitLog(context.Background(), 8) // offset = len("content\n") + done <- data + }() + + time.Sleep(10 * time.Millisecond) + sess.EnsureTrailingNewline() // no new byte, but signals end of turn + + select { + case data := <-done: + if data != nil { + t.Errorf("WaitLog = %q; want nil", data) + } + case <-time.After(time.Second): + t.Fatal("WaitLog did not unblock on end-of-turn") + } +} + func TestSessionStoreFileMode(t *testing.T) { if m, ok := store.SessionStoreFileMode("new"); !ok || m != 0666 { t.Errorf("SessionStoreFileMode(new) = %o, %v", m, ok)