15TurnLoop、取消与中断的实现

// TurnLoop is a push-based event loop for agent execution.
// Users push items via Push() and the loop processes them through the agent.
//
// Create with NewTurnLoop, then start with Run:
//
//  loop := NewTurnLoop(cfg)
//  // pass loop to other components, push initial items, etc.
//  loop.Run(ctx)
//
// # Permissive API
//
// All methods are valid on a not-yet-running loop:
//   - Push: items are buffered and will be processed once Run is called.
//   - Stop: sets the stopped flag; a subsequent Run will exit immediately.
//   - Wait: blocks until Run is called AND the loop exits. If Run is never
//     called, Wait blocks forever (this is a programming error, analogous
//     to reading from a channel that nobody writes to).
type TurnLoop[T any, M MessageType] struct {
  config TurnLoopConfig[T, M]

  buffer *turnBuffer[T]

  stopped int32
  started int32

  done chan struct{}

  result *TurnLoopExitState[T, M]

  runOnce sync.Once

  stopCtrl *stopController

  preemptCtrl *preemptController

  runErr error

  interruptedItems []T

  checkPointRunnerBytes []byte
  interruptContexts     []*InterruptCtx
  capturedCancelErr     *CancelError

  pendingResume *turnLoopPendingResume[T]

  loadCheckpointID string

  onAgentEvents func(ctx context.Context, tc *TurnContext[T, M], events *AsyncIterator[*TypedAgentEvent[M]]) error
// … 省略 5 行；完整声明 L879–931，点击上方「浏览完整文件」

// Run starts the loop's processing goroutine. It is non-blocking: the loop
// runs in the background and results are obtained via Wait.
//
// If CheckpointID is configured in TurnLoopConfig and a matching checkpoint
// exists in Store, the loop automatically resumes from that checkpoint.
// Otherwise it starts fresh with whatever items were Push()-ed.
//
// Calling Run more than once is a no-op: only the first call starts the loop.
func (l *TurnLoop[T, M]) Run(ctx context.Context) {
  l.start(ctx)
}

func (l *TurnLoop[T, M]) run(ctx context.Context) {
  defer l.cleanup(ctx)

  if err := l.tryLoadCheckpoint(ctx); err != nil {
    l.runErr = err
    return
  }

  // Monitor context cancellation: close the buffer so that a blocking
  // Receive() unblocks. The loop will then check ctx.Err() and exit.
  go func() {
    select {
    case <-ctx.Done():
      l.buffer.Close()
    case <-l.done:
    }
  }()

  for {
    if l.stopCtrl.isCommitted() {
      return
    }

    isResume := false
    var pr *turnLoopPendingResume[T]
    var items []T
    var pushBack []T

    if l.pendingResume != nil {
      isResume = true
      pr = l.pendingResume
      l.pendingResume = nil

      l.preemptCtrl.waitForPushes()
      pr.newItems = append(pr.newItems, l.buffer.TakeAll()...)

      pushBack = make([]T, 0, len(pr.interrupted)+len(pr.unhandled)+len(pr.newItems))
      pushBack = append(pushBack, pr.interrupted...)
      pushBack = append(pushBack, pr.unhandled...)
      pushBack = append(pushBack, pr.newItems...)
    } else {
      var first T
      var ok bool

      if idleFor := l.stopCtrl.idleDuration(); idleFor > 0 {
        l.buffer.ClearWakeup()
        idleTimer := time.NewTimer(idleFor)
        cancelIdle := make(chan struct{})
// … 省略 120 行；完整声明 L1570–1737，点击上方「浏览完整文件」

// Push adds an item to the loop's buffer for processing.
// This method is non-blocking and thread-safe.
// Returns false if the loop has stopped, true otherwise. If a preemptive push
// succeeds, the second return value is a channel that callers can wait on to
// confirm the preempt request has been resolved. Specifically:
//   - If Push observes a planning or active turn that is still the target when
//     the request resolves, the channel closes after TurnLoop attempts to submit
//     cancel for that target turn.
//   - If Push observes no target turn, the loop has not started, the preempt
//     subsystem is closed, or a delayed target is already gone, the channel
//     closes as a no-op resolution.
//
// If the loop has not been started yet (Run not called), items are buffered
// and will be processed once Run is called.
// After Wait() returns, failed pushes can be recovered via TurnLoopExitState.TakeLateItems().
// Once TakeLateItems() has been called, any subsequent push that would become a
// late item will panic instead of being silently dropped.
//
// Use WithPreempt() or WithPreemptTimeout() to atomically push an item and signal
// preemption of the current agent. This is useful for urgent items that should
// interrupt the current processing.
// The returned channel may be waited on if the caller needs to ensure the preempt
// signal has been observed.
//
// Use WithPreemptDelay() together with WithPreempt()/WithPreemptTimeout() to delay
// request resolution. Push returns immediately after the item is buffered, and
// the delayed request remains bound to the turn observed by Push.
func (l *TurnLoop[T, M]) Push(item T, opts ...PushOption[T, M]) (bool, <-chan struct{}) {
  cfg := &pushConfig[T, M]{}
  for _, opt := range opts {
    opt(cfg)
  }

  if cfg.pushStrategy != nil {
    return l.pushWithStrategy(item, cfg)
  }

  return l.pushWithConfig(item, cfg)
}

// Stop signals the loop to stop and returns immediately (non-blocking).
// Without options, the current agent turn runs to completion and the loop
// exits at the turn boundary without starting a new turn. ExitReason is nil.
//
// Use WithImmediate() to abort the running agent turn immediately.
// Use WithGraceful() to cancel at the nearest safe point with recursive
// propagation to nested agents.
// Use WithGracefulTimeout() for safe-point cancel with an escalation deadline.
// Use UntilIdleFor() to defer the stop until the loop has been continuously
// idle for a given duration; the loop shuts down automatically once the idle
// timer fires.
//
// This method may be called multiple times; subsequent calls update cancel options.
// A Stop() call without UntilIdleFor shuts down the loop immediately, even if
// a prior UntilIdleFor is still waiting.
// Call Wait() to block until the loop has fully exited and get the result.
//
// Stop may be called before Run. In that case, the stopped flag is set and
// a subsequent Run will exit the loop immediately.
//
// If the running agent does not support the WithCancel AgentRunOption,
// all cancel-related options (WithImmediate, WithGraceful, WithGracefulTimeout)
// degrade to "exit the loop on entering the next iteration" — the current
// agent turn runs to completion before the loop exits.
func (l *TurnLoop[T, M]) Stop(opts ...StopOption) {
  cfg := &stopConfig{}
  for _, opt := range opts {
    opt(cfg)
  }

  // UntilIdleFor is incompatible with cancel options (WithImmediate,
  // WithGraceful, WithGracefulTimeout) in the same call. Cancel opts only
  // make sense for an immediate or escalated stop; UntilIdleFor defers the
  // stop until idle, and must not impact a running agent. Drop them silently.
  if cfg.idleFor > 0 {
    cfg.agentCancelOpts = nil
  }

  decision := l.stopCtrl.requestStop(cfg)
  if decision.wakeIdle {
    l.buffer.Wakeup()
  }
  if decision.commit {
    l.finishStopCommit()
  }
}

// Wait blocks until the loop exits and returns the result.
// This method is safe to call from multiple goroutines.
// All callers will receive the same result.
//
// Wait blocks until Run is called AND the loop exits. If Run is
// never called, Wait blocks forever.
func (l *TurnLoop[T, M]) Wait() *TurnLoopExitState[T, M] {
  <-l.done
  return l.result
}

func (l *TurnLoop[T, M]) runAgentAndHandleEvents(
  ctx context.Context,
  agent TypedAgent[M],
  spec *turnRunSpec[T, M],
) error {
  l.interruptContexts = nil
  l.capturedCancelErr = nil
  l.checkPointRunnerBytes = nil

  var iter *AsyncIterator[*TypedAgentEvent[M]]

  runOpts, ms, err := l.setupBridgeStore(spec, spec.runOpts)
  if err != nil {
    l.preemptCtrl.abortPlanningTurn().ack()
    return err
  }
  store := l.config.Store
  cancelOpt, agentCancelFunc := WithCancel()
  runOpts = append(runOpts, cancelOpt)

  // For Run path the streaming mode comes from the input. For Resume path the
  // runner reads the streaming mode persisted in the checkpoint, so the value we
  // pass here is irrelevant.
  enableStreaming := false
  if spec.input != nil {
    enableStreaming = spec.input.EnableStreaming
  }
  runner := NewTypedRunner(TypedRunnerConfig[M]{
    EnableStreaming: enableStreaming,
    Agent:           agent,
    CheckPointStore: ms,
  })

  preemptDone := make(chan struct{})
  stoppedDone := make(chan struct{})

  tc := &TurnContext[T, M]{
    Loop:      l,
    Consumed:  spec.consumed,
    Preempted: preemptDone,
    Stopped:   stoppedDone,
    StopCause: l.stopCtrl.cause,
  }
  l.preemptCtrl.beginActiveTurn(ctx, tc)
  l.stopCtrl.beginActiveTurn()
  defer func() {
    l.stopCtrl.endActiveTurn()
    l.preemptCtrl.endActiveTurn().ack()
// … 省略 121 行；完整声明 L1809–1977，点击上方「浏览完整文件」

// CancelMode specifies when an agent should be canceled.
// Modes can be combined with bitwise OR to cancel at multiple safe-points.
// For example, CancelAfterChatModel | CancelAfterToolCalls cancels the agent
// after whichever safe-point is reached first.
type CancelMode int

// SafePoint describes at which boundary the agent may be cancelled.
// It is a bitmask: values can be combined with bitwise OR to accept multiple
// safe points (e.g. AfterToolCalls | AfterChatModel). Internally, SafePoint
// is translated to CancelMode via toCancelMode().
//
// SafePoint is used only in the preemption API (WithPreempt/WithPreemptTimeout).
// A key design constraint: preemption always targets a safe point — the user's
// intent is to cancel at a well-defined boundary, never to abort immediately.
// Immediate cancellation is only reachable as an automatic timeout escalation
// (via WithPreemptTimeout), not as a direct user choice. This is why SafePoint
// has no "immediate" value and why WithPreempt requires a non-zero SafePoint
// (panics otherwise).
type SafePoint int

// WithCancel creates an AgentRunOption that enables cancellation for an agent run.
// It returns the option to pass to Run/Resume and a cancel function.
// Cancel options (mode, timeout) are passed to the returned AgentCancelFunc at call time.
func WithCancel() (AgentRunOption, AgentCancelFunc) {
  cc := newCancelContext()
  opt := WrapImplSpecificOptFn(func(o *options) {
    o.cancelCtx = cc
  })
  cancelFn := cc.buildCancelFunc()
  return opt, cancelFn
}

func (cc *cancelContext) triggerCancel(mode CancelMode) {
  cc.setMode(mode)
  if atomic.CompareAndSwapInt32(&cc.state, stateRunning, stateCancelling) {
    close(cc.cancelChan)
  }
}

type cancelContext struct {
  mode int32 // atomic, CancelMode

  cancelChan    chan struct{} // closed when cancel is requested (all modes, not just safe-point)
  immediateChan chan struct{} // closed when an immediate graph interrupt fires
  doneChan      chan struct{} // closed when execution completes (by any mark* method)
  doneOnce      sync.Once     // ensures doneChan is closed exactly once

  state            int32 // stateRunning, stateCancelling, stateDone, stateCancelHandled
  interruptSent    int32 // interruptNotSent, interruptImmediate
  escalated        int32 // 1 if escalated from safe-point to immediate
  timeoutEscalated int32 // 1 if escalation was triggered by timeout
  startedMode      int32 // atomic, mode when state transitioned to cancelling
  deadlineUnixNano int64 // atomic, 0 means no deadline

  recursive     int32         // atomic; 1 if cancel should propagate into AgentTool internal agents
  recursiveChan chan struct{} // closed when recursive transitions from 0 to 1

  root   bool           // true for the original cancelContext created by WithCancel(); false for AgentTool internal agents
  parent *cancelContext // non-nil for AgentTool internal agents; used to propagate AgentTool boundary markers upward

  agentToolDescendant int32 // atomic; 1 once an AgentTool runs under this cancel context

  cancelMu      sync.Mutex
  timeoutOnce   sync.Once
  timeoutNotify chan struct{}

  mu                  sync.Mutex
  graphInterruptFuncs []func(...compose.GraphInterruptOption)
}

// shouldCancel returns true if a cancel has been requested (cancelChan is closed).
func (cc *cancelContext) shouldCancel() bool {
  if cc == nil {
    return false
  }
  select {
  case <-cc.cancelChan:
    return true
  default:
    return false
  }
}

// isImmediateCancelled returns true if an immediate graph interrupt has been
// fired (CancelImmediate or timeout escalation). This is stronger than
// shouldCancel: it means the compose graph is being torn down right now and
// orphaned goroutines should not attempt to send events.
func (cc *cancelContext) isImmediateCancelled() bool {
  if cc == nil {
    return false
  }
  select {
  case <-cc.immediateChan:
    return true
  default:
    return false
  }
}

// Interrupt creates a basic interrupt action.
// This is used when an agent needs to pause its execution to request external input or intervention,
// but does not need to save any internal state to be restored upon resumption.
// The `info` parameter is user-facing data that describes the reason for the interrupt.
func Interrupt(ctx context.Context, info any) *AgentEvent {
  return TypedInterrupt[*schema.Message](ctx, info)
}

// StatefulInterrupt creates an interrupt action that also saves the agent's internal state.
// This is used when an agent has internal state that must be restored for it to continue correctly.
// The `info` parameter is user-facing data describing the interrupt.
// The `state` parameter is the agent's internal state object, which will be serialized and stored.
func StatefulInterrupt(ctx context.Context, info any, state any) *AgentEvent {
  return TypedStatefulInterrupt[*schema.Message](ctx, info, state)
}

// TypedInterrupt creates a typed interrupt event that pauses execution to request external input.
// It is the generic counterpart of Interrupt; see Interrupt for full documentation.
func TypedInterrupt[M MessageType](ctx context.Context, info any) *TypedAgentEvent[M] {
  var rp []RunStep
  rCtx := getRunCtx(ctx)
  if rCtx != nil {
    rp = rCtx.RunPath
  }

  is, err := core.Interrupt(ctx, info, nil, nil,
    core.WithLayerPayload(rp))
  if err != nil {
    return &TypedAgentEvent[M]{Err: err}
  }

  contexts := core.ToInterruptContexts(is, allowedAddressSegmentTypes)

  return &TypedAgentEvent[M]{
    Action: &AgentAction{
      Interrupted: &InterruptInfo{
        InterruptContexts: contexts,
      },
      internalInterrupted: is,
    },
  }
}

// CompositeInterrupt creates an interrupt event that aggregates sub-interrupt signals.
func CompositeInterrupt(ctx context.Context, info any, state any,
  subInterruptSignals ...*InterruptSignal) *AgentEvent {
  return TypedCompositeInterrupt[*schema.Message](ctx, info, state, subInterruptSignals...)
}

// TypedCompositeInterrupt creates a typed interrupt event that aggregates sub-interrupt signals.
// It is the generic counterpart of CompositeInterrupt; see CompositeInterrupt for full documentation.
func TypedCompositeInterrupt[M MessageType](ctx context.Context, info any, state any,
  subInterruptSignals ...*InterruptSignal) *TypedAgentEvent[M] {
  var rp []RunStep
  rCtx := getRunCtx(ctx)
  if rCtx != nil {
    rp = rCtx.RunPath
  }

  is, err := core.Interrupt(ctx, info, state, subInterruptSignals,
    core.WithLayerPayload(rp))
  if err != nil {
    return &TypedAgentEvent[M]{Err: err}
  }

  contexts := core.ToInterruptContexts(is, allowedAddressSegmentTypes)

  return &TypedAgentEvent[M]{
    Action: &AgentAction{
      Interrupted: &InterruptInfo{
        InterruptContexts: contexts,
      },
      internalInterrupted: is,
    },
  }
}

type serialization struct {
    RunCtx              *runContext             // 整个运行环境(第 13 章)
    Info                *InterruptInfo          // 兼容旧版
    EnableStreaming     bool
    InterruptID2Address map[string]Address      // 中断 id → 节点地址
    InterruptID2State   map[string]core.InterruptState // 中断 id → 状态
}

// ResumeInfo holds all the information necessary to resume an interrupted agent execution.
// It is created by the framework and passed to an agent's Resume method.
type ResumeInfo struct {
  // EnableStreaming indicates whether the original execution was in streaming mode.
  EnableStreaming bool

  // Deprecated: use InterruptContexts from the embedded InterruptInfo for user-facing details,
  // and GetInterruptState for internal state retrieval.
  *InterruptInfo

  WasInterrupted bool
  InterruptState any
  IsResumeTarget bool
  ResumeData     any
}

type bridgeStore struct {
  mu   sync.Mutex
  data map[string][]byte
}

func (m *bridgeStore) Get(_ context.Context, key string) ([]byte, bool, error) {
  m.mu.Lock()
  defer m.mu.Unlock()
  if v, ok := m.data[key]; ok {
    return v, true, nil
  }
  return nil, false, nil
}

finalizeCheckpoint := func() error {
    if store != nil && ms != nil {
        data, ok, err := ms.Get(ctx, bridgeCheckpointID) // 从 bridgeStore 取回底层 checkpoint
        // ...
        if ok {
            l.checkPointRunnerBytes = append([]byte{}, data...) // 拼进 TurnLoop 的状态
        }
    }
    return nil
}

func newResumeBridgeStore(checkPointID string, data []byte) *bridgeStore {
  return &bridgeStore{
    data: map[string][]byte{checkPointID: data},
  }
}

func newBridgeStore() *bridgeStore {
  return &bridgeStore{data: make(map[string][]byte)}
}