fix(pty): consume drained session output

agent/inbox/inbox.go

@@ -19,6 +19,7 @@ type Inbox interface {
 	Closed() bool
 	Len() int
 	Wait(ctx context.Context) bool
+	WaitWhileActive(ctx context.Context) bool
 	RegisterProducer(name string) *ProducerHandle
 	ActiveProducers() int
 }
@@ -156,13 +157,21 @@ func (b *Buffered) Len() int {
 }
 
 func (b *Buffered) Wait(ctx context.Context) bool {
+	return b.wait(ctx, false)
+}
+
+func (b *Buffered) WaitWhileActive(ctx context.Context) bool {
+	return b.wait(ctx, true)
+}
+
+func (b *Buffered) wait(ctx context.Context, stopWhenIdle bool) bool {
 	for {
 		b.mu.Lock()
 		if len(b.buf) > 0 {
 			b.mu.Unlock()
 			return true
 		}
-		if b.closed {
+		if b.closed || (stopWhenIdle && len(b.producers) == 0) {
 			b.mu.Unlock()
 			return false
 		}

agent/inbox/inbox_test.go

@@ -1,8 +1,10 @@
 package inbox
 
 import (
+	"context"
 	"sync"
 	"testing"
+	"time"
 )
 
 func TestBufferedPushDrain(t *testing.T) {
@@ -173,3 +175,19 @@ func TestProducerRegistration(t *testing.T) {
 		t.Fatalf("expected 0 producers, got %d", b.ActiveProducers())
 	}
 }
+
+func TestBufferedWaitWhileActiveReturnsWhenProducersFinish(t *testing.T) {
+	b := NewBuffered(1)
+	producer := b.RegisterProducer("task")
+	ctx, cancel := context.WithTimeout(context.Background(), time.Second)
+	defer cancel()
+
+	go producer.Done()
+
+	if b.WaitWhileActive(ctx) {
+		t.Fatal("WaitWhileActive() reported a message after the producer finished")
+	}
+	if err := ctx.Err(); err != nil {
+		t.Fatalf("WaitWhileActive() did not return when the producer finished: %v", err)
+	}
+}

agent/loop.go

@@ -201,13 +201,10 @@ func runLoop(ctx context.Context, cfg Config) (*Result, error) {
 				continue
 			}
 
-			alive := (cfg.LoopScheduler != nil && cfg.LoopScheduler.Active() > 0) ||
-				(ib != nil && ib.ActiveProducers() > 0)
-
-			if alive && ib != nil && !ib.Closed() {
-				cfg.Logger.Debugf("[turn %d] waiting for inbox (loops=%d producers=%d)",
-					turn, schedulerActive(cfg.LoopScheduler), ib.ActiveProducers())
-				hasMessage := ib.Wait(ctx)
+			if ib != nil && !ib.Closed() {
+				cfg.Logger.Debugf("[turn %d] waiting for inbox (producers=%d)",
+					turn, ib.ActiveProducers())
+				hasMessage := ib.WaitWhileActive(ctx)
 				if hasMessage {
 					continue
 				}
@@ -713,13 +710,6 @@ func logUsage(logger telemetry.Logger, usage *aop.TokenUsage) {
 	}
 }
 
-func schedulerActive(s *LoopScheduler) int {
-	if s == nil {
-		return 0
-	}
-	return s.Active()
-}
-
 // messageBuilder accumulates streamed deltas into one assistant message.
 type messageBuilder struct {
 	role      string

agent/loop_scheduler.go

@@ -56,6 +56,7 @@ type LoopInfo struct {
 type LoopScheduler struct {
 	mu          sync.Mutex
 	loops       map[string]*loopState
+	ctx         context.Context
 	inbox       inbox.Inbox
 	log         telemetry.Logger
 	minInterval time.Duration
@@ -64,18 +65,23 @@ type LoopScheduler struct {
 type loopState struct {
 	entry     LoopEntry
 	cancel    context.CancelFunc
+	producer  *inbox.ProducerHandle
 	fireCount int
 	lastFired time.Time
 }
 
 const DefaultMinLoopInterval = 10 * time.Second
 
-func NewLoopScheduler(ib inbox.Inbox, logger telemetry.Logger) *LoopScheduler {
+func NewLoopScheduler(ctx context.Context, ib inbox.Inbox, logger telemetry.Logger) *LoopScheduler {
+	if ctx == nil {
+		ctx = context.Background()
+	}
 	if logger == nil {
 		logger = telemetry.NopLogger()
 	}
 	return &LoopScheduler{
 		loops:       make(map[string]*loopState),
+		ctx:         ctx,
 		inbox:       ib,
 		log:         logger,
 		minInterval: DefaultMinLoopInterval,
@@ -94,7 +100,7 @@ func (s *LoopScheduler) SetLogger(logger telemetry.Logger) {
 	s.mu.Unlock()
 }
 
-func (s *LoopScheduler) Add(ctx context.Context, entry LoopEntry) (string, error) {
+func (s *LoopScheduler) Add(entry LoopEntry) (string, error) {
 	if strings.TrimSpace(entry.Prompt) == "" {
 		return "", fmt.Errorf("prompt is required")
 	}
@@ -116,8 +122,12 @@ func (s *LoopScheduler) Add(ctx context.Context, entry LoopEntry) (string, error
 		s.mu.Unlock()
 		return "", fmt.Errorf("loop %q already exists", entry.Name)
 	}
-	loopCtx, cancel := context.WithCancel(ctx)
-	state := &loopState{entry: entry, cancel: cancel}
+	loopCtx, cancel := context.WithCancel(s.ctx)
+	state := &loopState{
+		entry:    entry,
+		cancel:   cancel,
+		producer: s.inbox.RegisterProducer("loop:" + entry.Name),
+	}
 	s.loops[entry.Name] = state
 	s.mu.Unlock()
 
@@ -136,6 +146,15 @@ func autoName(prompt string) string {
 }
 
 func (s *LoopScheduler) run(ctx context.Context, state *loopState) {
+	defer func() {
+		state.producer.Done()
+		s.mu.Lock()
+		if s.loops[state.entry.Name] == state {
+			delete(s.loops, state.entry.Name)
+		}
+		s.mu.Unlock()
+	}()
+
 	if state.entry.Cron != nil {
 		s.runCron(ctx, state)
 	} else {

agent/loop_scheduler_test.go

@@ -0,0 +1,37 @@
+package agent
+
+import (
+	"context"
+	"testing"
+	"time"
+
+	"github.com/chainreactors/aiscan/agent/inbox"
+)
+
+func TestLoopSchedulerProducerLifecycle(t *testing.T) {
+	ib := inbox.NewBuffered(1)
+	scheduler := NewLoopScheduler(context.Background(), ib, nil)
+
+	name, err := scheduler.Add(LoopEntry{
+		Name:     "test-loop",
+		Prompt:   "check progress",
+		Interval: time.Hour,
+	})
+	if err != nil {
+		t.Fatalf("Add() error = %v", err)
+	}
+	if got := ib.ActiveProducers(); got != 1 {
+		t.Fatalf("active producers after Add() = %d, want 1", got)
+	}
+
+	if err := scheduler.Remove(name); err != nil {
+		t.Fatalf("Remove() error = %v", err)
+	}
+	deadline := time.Now().Add(time.Second)
+	for ib.ActiveProducers() != 0 && time.Now().Before(deadline) {
+		time.Sleep(time.Millisecond)
+	}
+	if got := ib.ActiveProducers(); got != 0 {
+		t.Fatalf("active producers after Remove() = %d, want 0", got)
+	}
+}

go.mod

@@ -33,7 +33,7 @@ require (
 	github.com/chainreactors/utils v0.0.0-20260711153742-f3d210a5fa9d
 	github.com/chainreactors/utils/mitmproxy v0.0.0-20260818093021-b0af431aff73
 	github.com/chainreactors/utils/parsers v0.0.3
-	github.com/chainreactors/utils/pty v0.0.0-20260722180147-5b1816060721
+	github.com/chainreactors/utils/pty v0.0.0-20260819053645-5ed8693f0059
 	github.com/chainreactors/zombie v1.3.1-0.20260809133033-0d0df6fa50f5
 	github.com/charmbracelet/bubbles v1.0.0
 	github.com/charmbracelet/glamour v0.8.0

go.sum

@@ -263,8 +263,8 @@ github.com/chainreactors/utils/mitmproxy v0.0.0-20260818093021-b0af431aff73 h1:i
 github.com/chainreactors/utils/mitmproxy v0.0.0-20260818093021-b0af431aff73/go.mod h1:2O3/Vw66VnbzhwsHGFJ2Ge98RuSh6XzMMFGZmMmlZ9M=
 github.com/chainreactors/utils/parsers v0.0.3 h1:3ld7xG5TSvzikVOCkQHjqjHO3otjODwSwHQDkMKbu5o=
 github.com/chainreactors/utils/parsers v0.0.3/go.mod h1:bE/znJWt08n9QOORWsWu0ggB8GWfOg3+dfUMMITmwV4=
-github.com/chainreactors/utils/pty v0.0.0-20260722180147-5b1816060721 h1:gxkedbTvFEFTtel7XJEPMVh1iznfD+91woPkGBXZMNk=
-github.com/chainreactors/utils/pty v0.0.0-20260722180147-5b1816060721/go.mod h1:RW1v+8hFMeO9+TJyQ1iIx9Ea37s+B7BaDf9fyJ0OEC4=
+github.com/chainreactors/utils/pty v0.0.0-20260819053645-5ed8693f0059 h1:jnBzt8QOl9ekFKR2sBQsZdaeYUJlpA9PAky/KYmSwjg=
+github.com/chainreactors/utils/pty v0.0.0-20260819053645-5ed8693f0059/go.mod h1:RW1v+8hFMeO9+TJyQ1iIx9Ea37s+B7BaDf9fyJ0OEC4=
 github.com/chainreactors/words v0.0.0-20260520145736-270600e60fb4 h1:lvnDYEkatmZFHP5i321qQXK9L4vKRfso/uUfr5tOeC8=
 github.com/chainreactors/words v0.0.0-20260520145736-270600e60fb4/go.mod h1:zfz367PUmyaX6oAqV9SktVqyRXKlEh0sel9Wsq9dd2c=
 github.com/chainreactors/zombie v1.3.1-0.20260809133033-0d0df6fa50f5 h1:GB3a4+i5Yb1yO8eg9bTXtQUcttHOEs9KF75yq2a9EU4=

pkg/runner/loop_command.go

@@ -72,11 +72,11 @@ func (c *loopCommand) Run(ctx context.Context, execution *commands.Execution) (a
 		_, _ = fmt.Fprint(output, "All loops stopped.\n")
 		return nil, nil
 	default:
-		return nil, c.create(ctx, scheduler, output, args)
+		return nil, c.create(scheduler, output, args)
 	}
 }
 
-func (c *loopCommand) create(ctx context.Context, scheduler *agent.LoopScheduler, output io.Writer, args []string) error {
+func (c *loopCommand) create(scheduler *agent.LoopScheduler, output io.Writer, args []string) error {
 	if len(args) < 2 {
 		return fmt.Errorf("usage: loop <schedule> <prompt>")
 	}
@@ -98,7 +98,7 @@ func (c *loopCommand) create(ctx context.Context, scheduler *agent.LoopScheduler
 		return fmt.Errorf("usage: loop <schedule> <prompt>")
 	}
 
-	name, err := scheduler.Add(ctx, entry)
+	name, err := scheduler.Add(entry)
 	if err != nil {
 		return err
 	}

pkg/runner/runtime_session.go

@@ -526,6 +526,9 @@ func (m *sessionMailbox) Close()                        { m.base.Close() }
 func (m *sessionMailbox) Closed() bool                  { return m.base.Closed() }
 func (m *sessionMailbox) Len() int                      { return m.base.Len() }
 func (m *sessionMailbox) Wait(ctx context.Context) bool { return m.base.Wait(ctx) }
+func (m *sessionMailbox) WaitWhileActive(ctx context.Context) bool {
+	return m.base.WaitWhileActive(ctx)
+}
 func (m *sessionMailbox) RegisterProducer(name string) *inboxpkg.ProducerHandle {
 	return m.base.RegisterProducer(name)
 }
@@ -596,7 +599,7 @@ func (rt *AgentRuntime) OpenSession(ctx context.Context, options SessionOptions)
 	sessionCtx, cancel := context.WithCancel(ctx)
 	baseInbox := inboxpkg.NewBuffered(agent.DefaultInboxCapacity)
 	mailbox := &sessionMailbox{base: baseInbox}
-	scheduler := agent.NewLoopScheduler(mailbox, rt.config.Logger)
+	scheduler := agent.NewLoopScheduler(sessionCtx, mailbox, rt.config.Logger)
 	agentCfg := rt.config.
 		WithSystemPrompt(rt.systemPrompt).
 		WithStream(true).
@@ -628,7 +631,7 @@ func (rt *AgentRuntime) OpenSession(ctx context.Context, options SessionOptions)
 	rt.mu.Unlock()
 
 	if logicalID == MainREPLName && rt.option != nil && rt.option.Heartbeat > 0 {
-		_, _ = scheduler.Add(sessionCtx, agent.LoopEntry{
+		_, _ = scheduler.Add(agent.LoopEntry{
 			Name: "heartbeat", Interval: time.Duration(rt.option.Heartbeat) * time.Minute,
 			Mode:   agent.ModeInbox,
 			Prompt: "Heartbeat: review current context, check on any running sessions, and decide if action is needed.",