fix: 🐛 Keep in-flight pipeline running

cmd/agentop/edit/edit.go

@@ -174,13 +174,9 @@ type PolledMsg struct {
 }
 
 // ClusterPollDeadline bounds a ConfigMap edit's wait: the worst-case kubelet
-// sync (~60s) plus the framework's drain window (30s) plus jitter.
-//
-// The drain does NOT actually delay the verdict — the reloader records success
-// synchronously right after the Holder swap and drains the old pipelines in a
-// goroutine — so this is generous rather than tight. Kept at 120s regardless,
-// because for a cluster the kubelet sync is the real variable and it is not
-// worth shaving.
+// sync (~60s) plus jitter, doubled. The verdict does not wait on the drain —
+// the reloader records success right after the Holder swap and drains the old
+// pipelines in a goroutine.
 const ClusterPollDeadline = 120 * time.Second
 
 // LocalPollDeadline bounds a local file edit's wait. The proxy is already

core/listener/forwardproxy/websocket.go

@@ -329,11 +329,9 @@ func (ws *wsSocket) wsEnd(side string) {
 type wsTurn struct {
 	pctx *pipeline.Context
 	// p is the pipeline the turn began on, and every later step of the turn runs on
-	// it. A hot reload swaps OutboundPipeline for new requests and keeps the old one
-	// running for its drain window so requests already in flight finish where they
-	// started; an HTTPS turn does that by lasting seconds. A socket outlives the
-	// window, so a turn looking the pipeline up again at each step would end on one
-	// whose parser never saw it begin — no tokens, no cost.
+	// it, and the turn holds it until its Finishers run. A hot reload swaps
+	// OutboundPipeline for new requests, so a turn looking the pipeline up again at
+	// each step would end on one whose parser never saw it begin — no tokens, no cost.
 	p *pipeline.Pipeline
 	// prewarm is a generate:false response.create: run through the request phase, not
 	// recorded, and its answer not parsed.
@@ -611,17 +609,18 @@ func (ws *wsSocket) clientMessage(payload []byte) (replace []byte, refuse bool)
 	}
 	body, generate := responseCreateBody(payload)
 
-	p := ws.s.OutboundPipeline.Load()
+	pctx := ws.newTurnContext(body)
+	p := ws.s.OutboundPipeline.Pin(pctx)
 	if why := pipelineCannotCarry(p); why != "" {
 		// A reload since the handshake made this socket one the proxy would have
 		// declined. End it with "try again later": Codex reconnects, and the new
 		// handshake is answered 426 and falls back to HTTPS, where the plugin applies.
 		slog.Info("forward-proxy: ending a carried WebSocket after a reload", "host", ws.r.Host, "reason", why)
+		p.RunFinish(ws.ctx(), pctx, pipeline.Outcome{})
 		ws.closeBoth(wsframe.CloseTryAgainLater, "the proxy's pipeline changed; reconnect")
 		return nil, true
 	}
 
-	pctx := ws.newTurnContext(body)
 	t := &wsTurn{pctx: pctx, p: p, prewarm: !generate}
 	action := p.Run(ws.ctx(), pctx)
 	if action.Type == pipeline.Reject {

core/listener/forwardproxy/websocket_review_test.go

@@ -155,6 +155,9 @@ func TestResponsesWebSocket_ReloadAddingAnUncarriablePluginEndsTheSocket(t *test
 	if n != 0 {
 		t.Errorf("origin received %d response.create; the turn must not be sent without its plugin", n)
 	}
+	if n := fresh.Held(); n != 0 {
+		t.Errorf("the refused turn still holds the pipeline %d times", n)
+	}
 }
 
 // A turn over the 32 MiB a request body may be ends the socket with 1009, as HTTPS

core/listener/forwardproxy/websocket_test.go

@@ -1012,6 +1012,7 @@ func TestResponsesWebSocket_ReloadMidTurnFinishesOnTheTurnsPipeline(t *testing.T
 		_, _ = fr.ReadFrame()
 	})
 	c := h.dial(t, nil)
+	handshake := h.holder.Load().Held()
 	c.send(responseCreate("hello"))
 	<-midTurn
 
@@ -1022,7 +1023,11 @@ func TestResponsesWebSocket_ReloadMidTurnFinishesOnTheTurnsPipeline(t *testing.T
 	if err != nil {
 		t.Fatal(err)
 	}
+	old := h.holder.Load()
 	h.holder.Store(fresh)
+	if n := old.Held(); n != handshake+1 {
+		t.Errorf("Held = %d mid-turn, want %d: the open turn must hold the pipeline it began on", n, handshake+1)
+	}
 	close(resume)
 	c.readUntil("response.completed")
 	c.close()

core/pipeline/context.go

@@ -420,8 +420,9 @@ type Context struct {
 	// stay open for minutes), and pctx.dispatched holds indexes into the
 	// OLD pipeline's plugin slice — indexing the new one with them panics
 	// when the new chain is shorter and otherwise runs hooks on plugins
-	// that never saw OnRequest. Direct *Pipeline calls (not via a Holder)
-	// never reload under a live request and leave this nil.
+	// that never saw OnRequest. Set by Holder.Run and Holder.Pin, which also
+	// hold the pipeline until RunFinish so a reload does not Stop it under
+	// the request.
 	pinned *Pipeline
 
 	// outcome is populated exactly once by Pipeline.RunFinish immediately

core/pipeline/holder.go

@@ -11,9 +11,8 @@ import (
 // reloader Store()s a newly-built pipeline under the hood, and the
 // next Load returns it. atomic.Pointer guarantees a Load-then-use
 // sequence sees a fully-initialized pipeline — a request that started
-// before a swap finishes against the original *Pipeline, and the
-// reloader is responsible for not Stop()ing the old pipeline until a
-// drain window expires.
+// before a swap finishes against the original *Pipeline, which the
+// reloader Stops once no request holds it (Pipeline.Drained).
 //
 // A Holder has exactly one long-lived value; its identity is the slot
 // listeners reference, not the pipeline at any given moment.
@@ -40,13 +39,13 @@ func (h *Holder) Load() *Pipeline { return h.p.Load() }
 // new pipeline; in-flight requests that already Loaded the old one
 // continue to run against it. The caller owns the lifecycle of the
 // replaced pipeline — read the return value of the previous Load and
-// call its Stop(ctx) after the intended drain window.
+// call its Stop(ctx) once Drained is closed.
 func (h *Holder) Store(p *Pipeline) { h.p.Store(p) }
 
-// Run is equivalent to h.Load().Run(ctx, pctx). Offered as a
+// Run is equivalent to h.Pin(pctx).Run(ctx, pctx). Offered as a
 // one-liner so listeners can keep their Pipeline-era call sites
 // unchanged when they migrate their field type from *Pipeline to
-// *Holder. One atomic Load per call.
+// *Holder.
 //
 // Run also pins the loaded pipeline on pctx: every later phase on the
 // same request (RunResponse, RunResponseFrame, RunFinish, Resend)
@@ -56,9 +55,20 @@ func (h *Holder) Store(p *Pipeline) { h.p.Store(p) }
 // the pin, a shorter new chain panics in RunFinish and an equal-length
 // one runs hooks on plugins that never saw the request.
 func (h *Holder) Run(ctx context.Context, pctx *Context) Action {
-	p := h.p.Load()
-	pctx.pinned = p
-	return p.Run(ctx, pctx)
+	return h.Pin(pctx).Run(ctx, pctx)
+}
+
+// Pin pins pctx to the current pipeline and returns it. It re-checks the
+// pipeline after holding it, so a reload that Stores in between is either
+// seen here or sees the hold: the reloader counts holds only after its Store.
+func (h *Holder) Pin(pctx *Context) *Pipeline {
+	for {
+		p := h.p.Load()
+		p.pin(pctx)
+		if h.p.Load() == p {
+			return p
+		}
+	}
 }
 
 // ForPctx returns the pipeline a pctx-carrying phase should run on: the

core/pipeline/holder_reload_test.go

@@ -2,6 +2,7 @@ package pipeline
 
 import (
 	"context"
+	"sync"
 	"testing"
 )
 
@@ -93,3 +94,105 @@ func TestHolder_CapabilityProbesFollowThePin(t *testing.T) {
 		t.Error("HasStreamingResponders() = true, want false: the Holder itself serves new requests")
 	}
 }
+
+// A request holds the pipeline it was pinned to once from Run until RunFinish, even
+// when pinned again or finished twice, Drained closes on its release, and a
+// finished request pinned again holds nothing.
+func TestPipeline_HeldFromRunUntilRunFinish(t *testing.T) {
+	old := mustPipeline(t, &holderTestPlugin{name: "a"})
+	h := NewHolder(old)
+
+	pctx := &Context{Direction: Outbound}
+	h.Run(context.Background(), pctx)
+	h.Pin(pctx)
+	h.Store(mustPipeline(t, &holderTestPlugin{name: "b"}))
+	drained := old.Drained()
+	if got := old.Held(); got != 1 {
+		t.Fatalf("Held = %d during the request, want 1", got)
+	}
+	select {
+	case <-drained:
+		t.Fatal("Drained closed while a request held the pipeline")
+	default:
+	}
+
+	h.RunFinish(context.Background(), pctx, Outcome{})
+	h.RunFinish(context.Background(), pctx, Outcome{})
+	if got := old.Held(); got != 0 {
+		t.Errorf("Held = %d after RunFinish, want 0", got)
+	}
+	select {
+	case <-drained:
+	default:
+		t.Error("Drained still open after the last request finished")
+	}
+	if p := h.Pin(pctx); p != h.Load() {
+		t.Error("a finished request pinned again stayed on the pipeline it finished on")
+	}
+	if got := h.Load().Held(); got != 0 {
+		t.Errorf("the new pipeline is held %d times by a request that had already finished", got)
+	}
+}
+
+// heldDuringFinish records how many requests held its pipeline while it finished.
+type heldDuringFinish struct {
+	holderTestPlugin
+	p    *Pipeline
+	held int64
+}
+
+func (f *heldDuringFinish) OnFinish(context.Context, *Context) { f.held = f.p.Held() }
+
+// A request holds its pipeline until its Finishers return, so a reload cannot Stop
+// a plugin that is still finishing.
+func TestPipeline_FinishersRunWhileTheRequestHoldsThePipeline(t *testing.T) {
+	f := &heldDuringFinish{holderTestPlugin: holderTestPlugin{name: "f"}}
+	f.p = mustPipeline(t, f)
+	h := NewHolder(f.p)
+	pctx := &Context{Direction: Outbound}
+	h.Run(context.Background(), pctx)
+	h.RunFinish(context.Background(), pctx, Outcome{})
+	if f.held != 1 {
+		t.Errorf("Held = %d inside OnFinish, want 1", f.held)
+	}
+}
+
+// A request the reloader's wait-for-zero misses must not run on a stopped pipeline:
+// Holder.Pin either holds the pipeline before the reloader counts it, or moves to
+// the one Stored in its place.
+func TestHolder_PinNeverLandsOnADrainedPipeline(t *testing.T) {
+	h := NewHolder(mustPipeline(t, &holderTestPlugin{name: "a"}))
+	var stopped sync.Map
+	done := make(chan struct{})
+	var wg sync.WaitGroup
+	for range 8 {
+		wg.Add(1)
+		go func() {
+			defer wg.Done()
+			for {
+				select {
+				case <-done:
+					return
+				default:
+				}
+				pctx := &Context{Direction: Outbound}
+				p := h.Pin(pctx)
+				if _, ok := stopped.Load(p); ok {
+					t.Error("a request pinned a pipeline the reloader had already drained")
+				}
+				if p.Held() == 0 {
+					t.Error("Pin returned a pipeline the request does not hold")
+				}
+				h.RunFinish(context.Background(), pctx, Outcome{})
+			}
+		}()
+	}
+	for range 500 {
+		old := h.Load()
+		h.Store(mustPipeline(t, &holderTestPlugin{name: "a"}))
+		<-old.Drained()
+		stopped.Store(old, true)
+	}
+	close(done)
+	wg.Wait()
+}

core/pipeline/pipeline.go

@@ -4,6 +4,8 @@ import (
 	"context"
 	"fmt"
 	"log/slog"
+	"sync"
+	"sync/atomic"
 	"time"
 )
 
@@ -32,6 +34,63 @@ type Pipeline struct {
 	// requestWriters[i] is plugins[i]'s WritesRequestBody, read once at New for the
 	// same reason. SetRequestModel is accepted only from a plugin that declares it.
 	requestWriters []bool
+	// held counts requests pinned to this pipeline whose RunFinish has not run.
+	held atomic.Int64
+	// drain is set by Drained and closed by the release that brings held to zero.
+	drain atomic.Pointer[drainSignal]
+}
+
+type drainSignal struct {
+	once sync.Once
+	ch   chan struct{}
+}
+
+func (d *drainSignal) close() { d.once.Do(func() { close(d.ch) }) }
+
+// pin makes p the pipeline every later phase of pctx's request runs on, and holds p
+// until that request's RunFinish. Listeners pin through Holder, which closes the race
+// with a reload's Store.
+func (p *Pipeline) pin(pctx *Context) {
+	if pctx.pinned == p {
+		return
+	}
+	if pctx.pinned != nil && !pctx.finished {
+		pctx.pinned.release()
+	}
+	pctx.pinned = p
+	if !pctx.finished {
+		p.held.Add(1)
+	}
+}
+
+func (p *Pipeline) release() {
+	if p.held.Add(-1) == 0 {
+		if d := p.drain.Load(); d != nil {
+			p.closeIfIdle(d)
+		}
+	}
+}
+
+// closeIfIdle closes d if nothing holds p now. The count is read again after d is
+// seen: a release that hit zero before a request pinned p must not close it.
+func (p *Pipeline) closeIfIdle(d *drainSignal) {
+	if p.held.Load() == 0 {
+		d.close()
+	}
+}
+
+// Held is the number of requests pinned to p whose RunFinish has not run.
+func (p *Pipeline) Held() int64 { return p.held.Load() }
+
+// Drained returns a channel closed once no request holds p. The reloader calls it
+// after Storing p's replacement, when no new request can pin p.
+func (p *Pipeline) Drained() <-chan struct{} {
+	d := &drainSignal{ch: make(chan struct{})}
+	if !p.drain.CompareAndSwap(nil, d) {
+		d = p.drain.Load()
+	}
+	p.closeIfIdle(d)
+	return d.ch
 }
 
 // Option configures pipeline construction.
@@ -598,6 +657,10 @@ func (p *Pipeline) RunFinish(ctx context.Context, pctx *Context, outcome Outcome
 		return
 	}
 	pctx.finished = true
+	if pctx.pinned == p {
+		// After the Finishers, which still run on p.
+		defer p.release()
+	}
 	if len(pctx.dispatched) == 0 {
 		return
 	}