From 16bf2115fa105e9a211974ebb5b396f66794ca55 Mon Sep 17 00:00:00 2001
From: Ollie Agent <ollie@example.com>
Date: Mon, 17 Aug 2026 17:57:14 +0200
Subject: [PATCH] Cancel active agent work during shutdown

---
 cmd/olliesrv/internal/agent/agent.go | 23 +++++++++++++++++--
 cmd/olliesrv/internal/agent/turn.go  | 34 ++++++++++++++++++++++++----
 2 files changed, 50 insertions(+), 7 deletions(-)

diff --git a/cmd/olliesrv/internal/agent/agent.go b/cmd/olliesrv/internal/agent/agent.go
index 0fdfec8..c2b29ec 100644
--- a/cmd/olliesrv/internal/agent/agent.go
+++ b/cmd/olliesrv/internal/agent/agent.go
@@ -33,6 +33,8 @@ type Agent struct {
 	newToolServer func() *toolclient.ToolsrvConn
 	newBackend    func(string) (backend.Backend, error)
 	currentAction atomic.Pointer[actionHandle]
+	closed        atomic.Bool
+	actionMu      sync.Mutex // serializes action registration with shutdown
 	warnedContext bool
 	resultCache   sync.Map
 
@@ -379,6 +381,8 @@ func (ag *Agent) IsRunning() bool {
 // Interrupt cancels the current in-progress agent turn.
 // Returns true if an action was running and was cancelled.
 func (ag *Agent) Interrupt(cause error) bool {
+	ag.actionMu.Lock()
+	defer ag.actionMu.Unlock()
 	if h := ag.currentAction.Load(); h != nil {
 		h.cancel(cause)
 		return true
@@ -505,6 +509,7 @@ type EventHandler func(Event)
 // actionHandle holds the cancel function for the current agent turn.
 type actionHandle struct {
 	cancel context.CancelCauseFunc
+	done   chan struct{}
 }
 
 // WatchField names supported by Agent.WaitChange.
@@ -532,9 +537,23 @@ func (ag *Agent) SetEnv(key, value string) {
 	ag.runtime.ToolServer.SetEnv(key, value)
 }
 
-// Close releases agent resources (dispatcher, execute server).
+// Close releases agent resources (dispatcher, execute server). It first prevents
+// new work, cancels the active turn, and waits for that turn to release its
+// references to the runtime before closing the tools connection.
 func (ag *Agent) Close() {
-	if ag.runtime.ToolServer != nil {
+	ag.closed.Store(true)
+
+	ag.actionMu.Lock()
+	h := ag.currentAction.Load()
+	if h != nil {
+		h.cancel(context.Canceled)
+	}
+	ag.actionMu.Unlock()
+	if h != nil {
+		<-h.done
+	}
+
+	if ag.runtime != nil && ag.runtime.ToolServer != nil {
 		ag.runtime.ToolServer.Close()
 	}
 }
diff --git a/cmd/olliesrv/internal/agent/turn.go b/cmd/olliesrv/internal/agent/turn.go
index 5699918..366043b 100644
--- a/cmd/olliesrv/internal/agent/turn.go
+++ b/cmd/olliesrv/internal/agent/turn.go
@@ -39,6 +39,9 @@ func (ag *Agent) Submit(ctx context.Context, input string) {
 		}
 	}()
 	ag.log.Debug("Agent.Submit() input_len=%d running=%v", len(input), ag.IsRunning())
+	if ag.closed.Load() {
+		return
+	}
 	if input == "" {
 		return
 	}
@@ -56,6 +59,9 @@ func (ag *Agent) Submit(ctx context.Context, input string) {
 	// with an executeTurn arriving via prompt.
 	ag.submitMu.Lock()
 	defer ag.submitMu.Unlock()
+	if ag.closed.Load() {
+		return
+	}
 
 	if ag.HandleCommand(ctx, input) {
 		return
@@ -115,11 +121,20 @@ func (ag *Agent) executeTurn(ctx context.Context, input string) string {
 	}
 
 	actCtx, actCancel := context.WithCancelCause(ctx)
-	handle := &actionHandle{cancel: actCancel}
+	turnDone := make(chan struct{})
+	handle := &actionHandle{cancel: actCancel, done: turnDone}
+	ag.actionMu.Lock()
+	if ag.closed.Load() {
+		ag.actionMu.Unlock()
+		actCancel(context.Canceled)
+		return ""
+	}
 	ag.currentAction.Store(handle)
+	ag.actionMu.Unlock()
 	defer func() {
 		if ag.currentAction.CompareAndSwap(handle, nil) {
 			actCancel(context.Canceled)
+			close(handle.done)
 		}
 	}()
 	ag.SetState("thinking")
@@ -197,7 +212,7 @@ func (ag *Agent) executeTurn(ctx context.Context, input string) string {
 		if compactLimit := ag.autoCompactLimit(actCtx); compactLimit > 0 && tokens >= compactLimit {
 			ag.emit(Event{Role: "info", Content: "auto-compacting context...\n"})
 			ag.SetState("compacting")
-			if _, err := ag.runCompact(context.WithoutCancel(actCtx), "auto"); err != nil {
+			if _, err := ag.runCompact(actCtx, "auto"); err != nil {
 				ag.log.Error("auto-compact failed: %v", err)
 			}
 			ag.SetState("thinking")
@@ -246,8 +261,10 @@ func (ag *Agent) executeTurn(ctx context.Context, input string) string {
 	)
 	for {
 		err = ag.run(actCtx)
-		actCancel(nil)
-		ag.currentAction.CompareAndSwap(handle, nil)
+		if ag.currentAction.CompareAndSwap(handle, nil) {
+			actCancel(nil)
+			close(handle.done)
+		}
 
 		if err == nil {
 			break
@@ -273,8 +290,15 @@ func (ag *Agent) executeTurn(ctx context.Context, input string) string {
 			ag.history.resetTurnAccumulators()
 			replyBuf.Reset()
 			actCtx, actCancel = context.WithCancelCause(ctx)
-			handle = &actionHandle{cancel: actCancel}
+			handle = &actionHandle{cancel: actCancel, done: make(chan struct{})}
+			ag.actionMu.Lock()
+			if ag.closed.Load() {
+				ag.actionMu.Unlock()
+				actCancel(context.Canceled)
+				break
+			}
 			ag.currentAction.Store(handle)
+			ag.actionMu.Unlock()
 			// Re-inject preamble — compaction strips it from history.
 			ag.history.injectPreamble(ag.runtime.PreambleString())
 			continue