diff --git a/agent/inbox/inbox.go b/agent/inbox/inbox.go index 582f1be8..2441bc5a 100644 --- a/agent/inbox/inbox.go +++ b/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 } diff --git a/agent/inbox/inbox_test.go b/agent/inbox/inbox_test.go index 407af2bb..4189a875 100644 --- a/agent/inbox/inbox_test.go +++ b/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) + } +} diff --git a/agent/loop.go b/agent/loop.go index 70e1d650..15b7da4f 100644 --- a/agent/loop.go +++ b/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 diff --git a/agent/loop_scheduler.go b/agent/loop_scheduler.go index cefead78..1ba85ddc 100644 --- a/agent/loop_scheduler.go +++ b/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 { diff --git a/agent/loop_scheduler_test.go b/agent/loop_scheduler_test.go new file mode 100644 index 00000000..30009ef0 --- /dev/null +++ b/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) + } +} diff --git a/go.mod b/go.mod index 30b686e8..987d1132 100644 --- a/go.mod +++ b/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 diff --git a/go.sum b/go.sum index 48468c79..4b895c3b 100644 --- a/go.sum +++ b/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= diff --git a/pkg/runner/loop_command.go b/pkg/runner/loop_command.go index 14617a4d..6d99ff12 100644 --- a/pkg/runner/loop_command.go +++ b/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 ") } @@ -98,7 +98,7 @@ func (c *loopCommand) create(ctx context.Context, scheduler *agent.LoopScheduler return fmt.Errorf("usage: loop ") } - name, err := scheduler.Add(ctx, entry) + name, err := scheduler.Add(entry) if err != nil { return err } diff --git a/pkg/runner/runtime_session.go b/pkg/runner/runtime_session.go index c55f89b0..3f1465d5 100644 --- a/pkg/runner/runtime_session.go +++ b/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.",