From 49cdcd70f275c858b0511a2e88ab30a48b157fa3 Mon Sep 17 00:00:00 2001 From: Claude Agent Date: Sat, 11 Apr 2026 18:26:42 +0000 Subject: fix: tie pool submissions to server lifecycle context Fix 1 (server.go): Replace context.Background() with s.ctx in handleAnswerQuestion, handleResumeTimedOutTask, and handleRunTask. Add a ctx field to Server (defaulting to context.Background()) and a SetContext method so the serve command can wire in the signal- cancellable lifecycle context. Fix 2 (serve.go): Call srv.SetContext(ctx) before StartHub so all pool submissions use the server's root context (already cancelled on SIGTERM/SIGINT). Pool.Shutdown and its wiring were already present. Co-Authored-By: Claude Sonnet 4.6 --- internal/api/server.go | 14 +++++++++++--- 1 file changed, 11 insertions(+), 3 deletions(-) (limited to 'internal/api') diff --git a/internal/api/server.go b/internal/api/server.go index 3bc4147..604f354 100644 --- a/internal/api/server.go +++ b/internal/api/server.go @@ -31,6 +31,7 @@ type questionStore interface { // Server provides the REST API and WebSocket endpoint for Claudomator. type Server struct { + ctx context.Context // server lifecycle context; used for pool submissions store *storage.DB logStore logStore // injectable for tests; defaults to store taskLogStore taskLogStore // injectable for tests; defaults to store @@ -63,6 +64,12 @@ func (s *Server) SetAPIToken(token string) { s.apiToken = token } +// SetContext replaces the server's lifecycle context used for pool submissions. +// Call this before StartHub to tie task submissions to the server's shutdown signal. +func (s *Server) SetContext(ctx context.Context) { + s.ctx = ctx +} + // SetNotifier configures a notifier that is called on every task completion. func (s *Server) SetNotifier(n notify.Notifier) { s.notifier = n @@ -85,6 +92,7 @@ func (s *Server) Pool() *executor.Pool { return s.pool } func NewServer(store *storage.DB, pool *executor.Pool, logger *slog.Logger, claudeBinPath, geminiBinPath string) *Server { wd, _ := os.Getwd() s := &Server{ + ctx: context.Background(), store: store, logStore: store, taskLogStore: store, @@ -344,7 +352,7 @@ func (s *Server) handleAnswerQuestion(w http.ResponseWriter, r *http.Request) { ResumeAnswer: input.Answer, SandboxDir: latest.SandboxDir, } - if err := s.pool.SubmitResume(context.Background(), tk, resumeExec); err != nil { + if err := s.pool.SubmitResume(s.ctx, tk, resumeExec); err != nil { writeJSON(w, http.StatusServiceUnavailable, map[string]string{"error": err.Error()}) return } @@ -389,7 +397,7 @@ func (s *Server) handleResumeTimedOutTask(w http.ResponseWriter, r *http.Request ResumeSessionID: latest.SessionID, ResumeAnswer: resumeMsg, } - if err := s.pool.SubmitResume(context.Background(), tk, resumeExec); err != nil { + if err := s.pool.SubmitResume(s.ctx, tk, resumeExec); err != nil { writeJSON(w, http.StatusServiceUnavailable, map[string]string{"error": err.Error()}) return } @@ -661,7 +669,7 @@ func (s *Server) handleRunTask(w http.ResponseWriter, r *http.Request) { // task isn't immediately re-cancelled by checkDepsReady. s.cascadeRetryDeps(r.Context(), originalTask) - if err := s.pool.Submit(context.Background(), t); err != nil { + if err := s.pool.Submit(s.ctx, t); err != nil { writeJSON(w, http.StatusServiceUnavailable, map[string]string{"error": fmt.Sprintf("executor pool: %v", err)}) return } -- cgit v1.2.3