From 2dedc0d20cbfd8f33be146635755293a19cdf357 Mon Sep 17 00:00:00 2001 From: weichen Date: Thu, 1 Oct 2026 10:11:19 +0000 Subject: [PATCH] Add an AgentSession with steer and follow-up queues Why: Steering hooks existed on the loop but nothing owned the queues, abort, or a model change during a run. Evidence: go test -count=1 ./runtime/ Impact: Run still emits the existing event names. message_start, tool_execution_update, and agent_settled are added beside them. Refs: feature/pi-headless N3 Co-authored-by: Cursor --- runtime/loop.go | 23 +++- runtime/session_agent.go | 200 ++++++++++++++++++++++++++++++++++ runtime/session_agent_test.go | 88 +++++++++++++++ 3 files changed, 309 insertions(+), 2 deletions(-) create mode 100644 runtime/session_agent.go create mode 100644 runtime/session_agent_test.go diff --git a/runtime/loop.go b/runtime/loop.go index 88a52a8..475152e 100644 --- a/runtime/loop.go +++ b/runtime/loop.go @@ -46,6 +46,10 @@ type Config struct { GetSteering func() []llm.Message GetFollowUp func() []llm.Message OnEvent func(Event) + // ThinkingLevel is forwarded on each provider request. + ThinkingLevel string + // Resolve, when set, overrides Model and ThinkingLevel at the start of each turn. + Resolve func() (model, thinking string) // Log 可选。nil 时丢弃日志，不改变循环结果。 Log *logsys.Logger // ExecutionID 只用于日志字段，不参与控制流。 @@ -117,6 +121,7 @@ func Run(ctx context.Context, cfg Config, user llm.Message) (Result, error) { return budgetStop(cfg, last, messages, usage, reason), nil } emit(cfg, Event{Type: "turn_start"}) + emit(cfg, Event{Type: "message_start"}) end, errEv := consume(ctx, cfg, messages) if errEv != nil { cfg.logger().Error("runtime turn failed", errEv, "retryable", end.Retryable) @@ -167,6 +172,7 @@ func Run(ctx context.Context, cfg Config, user llm.Message) (Result, error) { continue } emit(cfg, Event{Type: "agent_end", Message: last}) + emit(cfg, Event{Type: "agent_settled", Message: last}) emitUsage(cfg, usage) return Result{Text: last.Content, Messages: messages, StopReason: last.StopReason, Usage: usage}, nil } @@ -222,9 +228,21 @@ func consume(ctx context.Context, cfg Config, messages []llm.Message) (llm.Event decls = append(decls, llm.ToolDecl{Name: d.Name, Description: d.Description, Parameters: d.Parameters}) } } + model := cfg.Model + opts := llm.StreamOptions{ThinkingLevel: cfg.ThinkingLevel} + if cfg.Resolve != nil { + if m, th := cfg.Resolve(); m != "" || th != "" { + if m != "" { + model = m + } + if th != "" { + opts.ThinkingLevel = th + } + } + } var end llm.Event sawEnd := false - for ev := range cfg.Stream(ctx, cfg.Model, messages, decls, llm.StreamOptions{}) { + for ev := range cfg.Stream(ctx, model, messages, decls, opts) { switch ev.Kind { case llm.KindTextDelta: emit(cfg, Event{Type: "text_delta", Text: ev.Text}) @@ -271,7 +289,8 @@ func executeParallel(ctx context.Context, cfg Config, assistant llm.Message) []l } func runOne(ctx context.Context, cfg Config, call llm.ToolCall) llm.Message { - emit(cfg, Event{Type: "tool_execution_start", Tool: call.Name}) + emit(cfg, Event{Type: "tool_execution_start", Tool: call.Name, Detail: call.Arguments}) + emit(cfg, Event{Type: "tool_execution_update", Tool: call.Name, Detail: call.Arguments}) res := tools.Result{Content: "no tools configured", IsError: true} if cfg.Tools != nil { res = cfg.Tools.Execute(ctx, call.Name, json.RawMessage(call.Arguments)) diff --git a/runtime/session_agent.go b/runtime/session_agent.go new file mode 100644 index 0000000..1f07f21 --- /dev/null +++ b/runtime/session_agent.go @@ -0,0 +1,200 @@ +package runtime + +import ( + "context" + "sync" + + "agent-native/runtime/llm" + "agent-native/runtime/session" +) + +// Delivery selects how a queued steer or follow-up is handed to the loop. +type Delivery string + +const ( + DeliverOne Delivery = "one" + DeliverAll Delivery = "all" +) + +// AgentSession is the headless session object: queues, abort, and +// mid-run model changes around Run. +type AgentSession struct { + cfg Config + + mu sync.Mutex + steer []llm.Message + follow []llm.Message + steerMode Delivery + followMode Delivery + model string + thinking string + cancel context.CancelFunc + runID int +} + +// NewAgentSession copies cfg. Queue modes default to one-at-a-time. +func NewAgentSession(cfg Config) *AgentSession { + if cfg.Model == "" { + cfg.Model = "default" + } + return &AgentSession{ + cfg: cfg, + steerMode: DeliverOne, + followMode: DeliverOne, + model: cfg.Model, + thinking: cfg.ThinkingLevel, + } +} + +// SetQueueMode sets the delivery mode for steer and follow-up. +func (a *AgentSession) SetQueueMode(steer, follow Delivery) { + a.mu.Lock() + defer a.mu.Unlock() + if steer != "" { + a.steerMode = steer + } + if follow != "" { + a.followMode = follow + } +} + +// Steer queues a message to be delivered after the current tool batch. +func (a *AgentSession) Steer(msg llm.Message) { + a.mu.Lock() + a.steer = append(a.steer, msg) + a.mu.Unlock() + a.emitQueue("steer") +} + +// FollowUp queues a message for after the assistant stops. +func (a *AgentSession) FollowUp(msg llm.Message) { + a.mu.Lock() + a.follow = append(a.follow, msg) + a.mu.Unlock() + a.emitQueue("follow_up") +} + +// ClearQueue drops both queues. +func (a *AgentSession) ClearQueue() { + a.mu.Lock() + a.steer = nil + a.follow = nil + a.mu.Unlock() + a.emitQueue("clear") +} + +// Abort cancels the in-flight Prompt. +func (a *AgentSession) Abort() { + a.mu.Lock() + cancel := a.cancel + a.mu.Unlock() + if cancel != nil { + cancel() + } +} + +// SetModel changes the model and thinking level used on the next turn. +// A session store, when present, records the change. +func (a *AgentSession) SetModel(model, thinking string) error { + if thinking != "" { + norm, err := llm.NormalizeThinking(thinking) + if err != nil { + return err + } + thinking = norm + } + a.mu.Lock() + if model != "" { + a.model = model + } + if thinking != "" { + a.thinking = thinking + } + store := a.cfg.Session + a.mu.Unlock() + if store == nil { + return nil + } + if model != "" { + if _, err := store.AppendEntry(session.Entry{Type: session.TypeModelChange, Model: model}); err != nil { + return err + } + } + if thinking != "" { + if _, err := store.AppendEntry(session.Entry{Type: session.TypeThinkingLevelChange, ThinkingLevel: thinking}); err != nil { + return err + } + } + return nil +} + +// Prompt runs one user message until the agent settles. +func (a *AgentSession) Prompt(ctx context.Context, user llm.Message) (Result, error) { + ctx, cancel := context.WithCancel(ctx) + a.mu.Lock() + a.runID++ + id := a.runID + a.cancel = cancel + cfg := a.cfg + a.mu.Unlock() + defer func() { + cancel() + a.mu.Lock() + if a.runID == id { + a.cancel = nil + } + a.mu.Unlock() + }() + cfg.GetSteering = func() []llm.Message { return a.take("steer") } + cfg.GetFollowUp = func() []llm.Message { return a.take("follow") } + cfg.Resolve = func() (string, string) { + a.mu.Lock() + defer a.mu.Unlock() + return a.model, a.thinking + } + return Run(ctx, cfg, user) +} + +// Compact appends a compaction entry and emits compaction events. +func (a *AgentSession) Compact(keep int, summarize func([]llm.Message) (string, error)) (bool, error) { + if a.cfg.Session == nil { + return false, nil + } + emit(a.cfg, Event{Type: "compaction_start"}) + done, err := a.cfg.Session.Compact(keep, summarize) + text := "noop" + if done { + text = "compacted" + } + if err != nil { + text = err.Error() + } + emit(a.cfg, Event{Type: "compaction_end", Text: text}) + return done, err +} + +func (a *AgentSession) take(kind string) []llm.Message { + a.mu.Lock() + defer a.mu.Unlock() + q := &a.steer + mode := a.steerMode + if kind == "follow" { + q = &a.follow + mode = a.followMode + } + if len(*q) == 0 { + return nil + } + if mode == DeliverAll { + out := append([]llm.Message(nil), (*q)...) + *q = nil + return out + } + out := []llm.Message{(*q)[0]} + *q = (*q)[1:] + return out +} + +func (a *AgentSession) emitQueue(text string) { + emit(a.cfg, Event{Type: "queue_update", Text: text}) +} diff --git a/runtime/session_agent_test.go b/runtime/session_agent_test.go new file mode 100644 index 0000000..ed432a2 --- /dev/null +++ b/runtime/session_agent_test.go @@ -0,0 +1,88 @@ +package runtime + +import ( + "context" + "sync" + "testing" + + "agent-native/runtime/llm" +) + +func TestAgentSessionSteerFollowAndAbort(t *testing.T) { + var mu sync.Mutex + step := 0 + stream := llm.StreamFunc(func(ctx context.Context, model string, messages []llm.Message, td []llm.ToolDecl, opts llm.StreamOptions) <-chan llm.Event { + ch := make(chan llm.Event, 2) + go func() { + defer close(ch) + mu.Lock() + n := step + step++ + mu.Unlock() + if opts.ThinkingLevel != "low" && n == 0 { + t.Errorf("thinking level not forwarded: %s", opts.ThinkingLevel) + } + if n == 0 { + ch <- llm.Event{Kind: llm.KindMessageEnd, Message: llm.Message{ + Role: llm.RoleAssistant, StopReason: llm.StopToolCalls, + ToolCalls: []llm.ToolCall{{ID: "1", Name: "missing", Arguments: "{}"}}, + }} + return + } + ch <- llm.Event{Kind: llm.KindMessageEnd, Message: llm.Message{ + Role: llm.RoleAssistant, Content: "done", StopReason: llm.StopEnd, + }} + }() + return ch + }) + var types []string + sess := NewAgentSession(Config{ + Model: "m", + ThinkingLevel: "low", + Stream: stream, + MaxTurns: 6, + OnEvent: func(ev Event) { + types = append(types, ev.Type) + }, + }) + sess.SetQueueMode(DeliverOne, DeliverAll) + sess.Steer(llm.Message{Role: llm.RoleUser, Content: "steer-me"}) + sess.FollowUp(llm.Message{Role: llm.RoleUser, Content: "later"}) + res, err := sess.Prompt(context.Background(), llm.Message{Role: llm.RoleUser, Content: "go"}) + if err != nil { + t.Fatal(err) + } + if res.StopReason != llm.StopEnd { + t.Fatalf("stop %s err %v", res.StopReason, res.Err) + } + found := map[string]bool{} + for _, typ := range types { + found[typ] = true + } + for _, want := range []string{"message_start", "queue_update", "tool_execution_update", "agent_settled"} { + if !found[want] { + t.Fatalf("missing %s in %v", want, types) + } + } +} + +func TestClearQueueAndSetModel(t *testing.T) { + sess := NewAgentSession(Config{Model: "old", Stream: func(ctx context.Context, model string, messages []llm.Message, td []llm.ToolDecl, opts llm.StreamOptions) <-chan llm.Event { + ch := make(chan llm.Event, 1) + ch <- llm.Event{Kind: llm.KindMessageEnd, Message: llm.Message{Role: llm.RoleAssistant, Content: model, StopReason: llm.StopEnd}} + close(ch) + return ch + }}) + sess.Steer(llm.Message{Role: llm.RoleUser, Content: "drop"}) + sess.ClearQueue() + if err := sess.SetModel("new-model", "high"); err != nil { + t.Fatal(err) + } + res, err := sess.Prompt(context.Background(), llm.Message{Role: llm.RoleUser, Content: "hi"}) + if err != nil { + t.Fatal(err) + } + if res.Text != "new-model" { + t.Fatalf("model switch did not apply: %q", res.Text) + } +}