Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
62 changes: 51 additions & 11 deletions internal/runner/comfy_session.go
Original file line number Diff line number Diff line change
Expand Up @@ -5,6 +5,7 @@ package runner

import (
"context"
"encoding/json"
"fmt"
"time"

Expand All @@ -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 {
Expand All @@ -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
}
}
Expand Down Expand Up @@ -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()
Expand Down Expand Up @@ -146,24 +151,59 @@ 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)
if err != nil {
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
}
Expand Down
83 changes: 83 additions & 0 deletions internal/runner/comfy_startup_test.go
Original file line number Diff line number Diff line change
@@ -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)
}
})
}
}
2 changes: 1 addition & 1 deletion internal/runner/runner.go
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -223,7 +224,6 @@ func Run(o Options) int {
}
if o.ExitComfyUIIdle {
go func() {
var session comfySession
complete := false
for {
if !complete {
Expand Down
Loading