From 3989e60eb5e4492a11869110b5dd7d8f0678befa Mon Sep 17 00:00:00 2001
From: Bacongineer <55918611+Bacongineer@users.noreply.github.com>
Date: Wed, 21 Oct 2026 12:33:23 -0400
Subject: [PATCH] fix(runner): retain fast ComfyUI completions across startup

Goal: release a completed session even when its first history read is unavailable or follows a fast job.

How: capture the session boundary before starting the child and exclude terminal records from that grant from both initial and recovered history baselines.

Key details: preserve lifetime and undated history, reject future timestamps, and keep busy queues protected. Cover startup refusals, immediate history, successful/error/interrupted jobs, prior or unknown records, and busy queues.
---
 internal/runner/comfy_session.go      | 62 ++++++++++++++++----
 internal/runner/comfy_startup_test.go | 83 +++++++++++++++++++++++++++
 internal/runner/runner.go             |  2 +-
 3 files changed, 135 insertions(+), 12 deletions(-)
 create mode 100644 internal/runner/comfy_startup_test.go

diff --git a/internal/runner/comfy_session.go b/internal/runner/comfy_session.go
index 75f7feb..fced5c0 100644
--- a/internal/runner/comfy_session.go
+++ b/internal/runner/comfy_session.go
@@ -5,6 +5,7 @@ package runner
 
 import (
 	"context"
+	"encoding/json"
 	"fmt"
 	"time"
 
@@ -15,9 +16,10 @@ import (
 // comfySession retains the original grant and work evidence across refused
 // attempts. Failure timers belong to each wait, not to the session lifetime.
 type comfySession struct {
-	granted  time.Time
-	baseline map[string]bool
-	observed bool
+	granted     time.Time
+	baseline    map[string]bool
+	observed    bool
+	initialized bool
 }
 
 func (session *comfySession) wait(o Options, done <-chan struct{}, say func(string, ...any)) bool {
@@ -42,8 +44,11 @@ func (session *comfySession) wait(o Options, done <-chan struct{}, say func(stri
 	// that read belongs to this grant. History is used only by session release.
 	if session.granted.IsZero() {
 		session.granted = time.Now()
+	}
+	if !session.initialized {
+		session.initialized = true
 		if o.ExitComfyUIIdle {
-			if keys, err := comfyHistoryKeys(o.DrainComfyUI, o.APIKey); err == nil {
+			if keys, err := comfyHistoryKeys(o.DrainComfyUI, o.APIKey, session.granted); err == nil {
 				session.baseline = keys
 			}
 		}
@@ -77,7 +82,7 @@ func (session *comfySession) wait(o Options, done <-chan struct{}, say func(stri
 				say("ComfyUI queue empty; handing off")
 				return true
 			case session.baseline == nil:
-				session.baseline, err = comfyHistoryKeys(o.DrainComfyUI, o.APIKey)
+				session.baseline, err = comfyHistoryKeys(o.DrainComfyUI, o.APIKey, session.granted)
 				if err != nil {
 					if histBad.IsZero() {
 						histBad = time.Now()
@@ -146,9 +151,9 @@ func (session *comfySession) wait(o Options, done <-chan struct{}, say func(stri
 	}
 }
 
-// comfyHistoryKeys reads the newest completed-work key, enough to detect
-// whether this grant has completed a job without fetching the full history.
-func comfyHistoryKeys(base, key string) (map[string]bool, error) {
+// comfyHistoryKeys excludes timestamped terminal work from the baseline
+// when it belongs to the current grant.
+func comfyHistoryKeys(base, key string, since time.Time) (map[string]bool, error) {
 	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
 	defer cancel()
 	history, err := comfy.History(httpauth.WithBearer(ctx, key), base)
@@ -156,14 +161,49 @@ func comfyHistoryKeys(base, key string) (map[string]bool, error) {
 		return nil, err
 	}
 	keys := make(map[string]bool, len(history))
-	for k := range history {
-		keys[k] = true
+	for k, record := range history {
+		if since.IsZero() || !comfyCompletedAfter(record, since) {
+			keys[k] = true
+		}
 	}
 	return keys, nil
 }
 
+// comfyCompletedAfter accepts only terminal work recorded in this grant. Missing,
+// malformed, or future timestamps cannot prove that lifetime history is new work.
+func comfyCompletedAfter(raw json.RawMessage, since time.Time) bool {
+	var record struct {
+		Status struct {
+			Messages [][]json.RawMessage `json:"messages"`
+		} `json:"status"`
+	}
+	if json.Unmarshal(raw, &record) != nil {
+		return false
+	}
+	for _, message := range record.Status.Messages {
+		if len(message) != 2 {
+			continue
+		}
+		var event string
+		if json.Unmarshal(message[0], &event) != nil ||
+			(event != "execution_success" && event != "execution_error" && event != "execution_interrupted") {
+			continue
+		}
+		var detail struct {
+			Timestamp int64 `json:"timestamp"`
+		}
+		if json.Unmarshal(message[1], &detail) == nil {
+			completed := time.UnixMilli(detail.Timestamp)
+			if completed.After(since) && !completed.After(time.Now()) {
+				return true
+			}
+		}
+	}
+	return false
+}
+
 func comfyFirstNewJob(base, key string, baseline map[string]bool) (bool, error) {
-	keys, err := comfyHistoryKeys(base, key)
+	keys, err := comfyHistoryKeys(base, key, time.Time{})
 	if err != nil {
 		return false, err
 	}
diff --git a/internal/runner/comfy_startup_test.go b/internal/runner/comfy_startup_test.go
new file mode 100644
index 0000000..7bffd53
--- /dev/null
+++ b/internal/runner/comfy_startup_test.go
@@ -0,0 +1,83 @@
+// Copyright 2026 Trevin Tindall (T-Crypt)
+// SPDX-License-Identifier: Apache-2.0
+
+package runner
+
+import (
+	"fmt"
+	"net/http"
+	"net/http/httptest"
+	"strings"
+	"sync/atomic"
+	"testing"
+	"time"
+)
+
+func TestComfySessionRecoversCompletedJobFromUnavailableBaseline(t *testing.T) {
+	for _, tc := range []struct {
+		name      string
+		event     string
+		offset    time.Duration
+		undated   bool
+		busy      bool
+		initialOK bool
+		want      bool
+	}{
+		{name: "completed_during_startup", event: "execution_success", want: true},
+		{name: "completed_before_first_history_read", event: "execution_success", initialOK: true, want: true},
+		{name: "failed_job_is_finished_work", event: "execution_error", want: true},
+		{name: "interrupted_job_is_finished_work", event: "execution_interrupted", want: true},
+		{name: "prior_job_is_not_this_session", event: "execution_success", offset: -time.Hour},
+		{name: "future_clock_is_not_proof", event: "execution_success", offset: time.Hour},
+		{name: "undated_history_is_not_proof", undated: true},
+		{name: "started_only_is_not_completion", event: "execution_start"},
+		{name: "busy_queue_keeps_lease", event: "execution_success", busy: true},
+	} {
+		t.Run(tc.name, func(t *testing.T) {
+			var reads atomic.Int32
+			completed := time.Now().Add(tc.offset).UnixMilli()
+			srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+				switch r.URL.Path {
+				case "/queue":
+					if tc.busy {
+						fmt.Fprint(w, `{"queue_running":[[0,"active"]],"queue_pending":[]}`)
+					} else {
+						fmt.Fprint(w, `{"queue_running":[],"queue_pending":[]}`)
+					}
+				case "/history":
+					if reads.Add(1) == 1 && !tc.initialOK {
+						http.Error(w, "history not ready", http.StatusServiceUnavailable)
+					} else if tc.undated {
+						fmt.Fprint(w, `{"record":{}}`)
+					} else {
+						fmt.Fprintf(w, `{"record":{"status":{"messages":[[%q,{"timestamp":%d}]]}}}`, tc.event, completed)
+					}
+				default:
+					http.NotFound(w, r)
+				}
+			}))
+			defer srv.Close()
+			done := make(chan struct{})
+			defer func() {
+				select {
+				case <-done:
+				default:
+					close(done)
+				}
+			}()
+			session := comfySession{granted: time.Now().Add(-time.Second)}
+			say := func(format string, args ...any) {
+				message := fmt.Sprintf(format, args...)
+				if strings.Contains(message, "waiting for the first ComfyUI job") ||
+					strings.Contains(message, "waiting for ComfyUI to finish") {
+					close(done)
+				}
+			}
+			got := session.wait(Options{DrainComfyUI: srv.URL, ExitComfyUIIdle: true,
+				FailClosed: true, Grace: 10 * time.Millisecond, Retry: time.Millisecond}, done, say)
+			if got != tc.want {
+				t.Fatalf("session release = %v, want %v; baseline must distinguish completed session work from lifetime history", got, tc.want)
+			}
+		})
+	}
+}
diff --git a/internal/runner/runner.go b/internal/runner/runner.go
index eafb32a..8d40a5b 100644
--- a/internal/runner/runner.go
+++ b/internal/runner/runner.go
@@ -174,6 +174,7 @@ func Run(o Options) int {
 	}
 
 	cmd := command(o.Argv, o.Output)
+	session := comfySession{granted: time.Now()}
 	if err := cmd.Start(); err != nil {
 		say("start %s: %v", o.Argv[0], err)
 		return 127
@@ -223,7 +224,6 @@ func Run(o Options) int {
 	}
 	if o.ExitComfyUIIdle {
 		go func() {
-			var session comfySession
 			complete := false
 			for {
 				if !complete {