fix: Response completion uses an unbounded context for durable-ID lookup and can leave runs stuck in_progress forever

cmd/serve_response_runs.go

@@ -18,6 +18,8 @@ import (
 	"github.com/samsaffron/term-llm/internal/tools"
 )
 
+const responseCompletionDurableLookupTimeout = 5 * time.Second
+
 type responseRunEvent struct {
 	Sequence int64
 	Event    string
@@ -1757,6 +1759,25 @@ func (s *serveServer) streamResponseRunEvents(ctx context.Context, w http.Respon
 	}
 }
 
+func (s *serveServer) completedResponseIDForSessionWithin(ctx context.Context, sessionID, fallbackID string, timeout time.Duration) string {
+	if fallbackID == "" || sessionID == "" {
+		return fallbackID
+	}
+	if ctx == nil {
+		ctx = context.Background()
+	}
+	lookupCtx, cancel := context.WithTimeout(context.WithoutCancel(ctx), timeout)
+	defer cancel()
+	if durableID := s.latestDurableResponseIDForSession(lookupCtx, sessionID); durableID != "" {
+		return durableID
+	}
+	return fallbackID
+}
+
+func (s *serveServer) completedResponseIDForSession(ctx context.Context, sessionID, fallbackID string) string {
+	return s.completedResponseIDForSessionWithin(ctx, sessionID, fallbackID, responseCompletionDurableLookupTimeout)
+}
+
 func (s *serveServer) startResponseRun(runtime *serveRuntime, stateful bool, replaceHistory bool, inputMessages []llm.Message, llmReq llm.Request, sessionID string, options startResponseRunOptions) (*responseRun, error) {
 	mgr := s.ensureResponseRuns()
 
@@ -1910,11 +1931,7 @@ func (s *serveServer) startResponseRun(runtime *serveRuntime, stateful bool, rep
 		if options.resetResponseIDsOnSuccess {
 			s.unregisterSessionResponseIDs(sessionID)
 		}
-		durableID := s.latestDurableResponseIDForSession(context.Background(), sessionID)
-		completedID := respID
-		if durableID != "" {
-			completedID = durableID
-		}
+		completedID := s.completedResponseIDForSession(runCtx, sessionID, respID)
 		if completedID != respID {
 			s.registerResponseID(runtime, respID, sessionID)
 		}

cmd/serve_response_runs_test.go

@@ -11,6 +11,7 @@ import (
 	"time"
 
 	"github.com/samsaffron/term-llm/internal/llm"
+	"github.com/samsaffron/term-llm/internal/session"
 )
 
 func TestEncodeTextDeltaPayloadMatchesJSONMarshalForInvalidUTF8(t *testing.T) {
@@ -32,6 +33,75 @@ func TestEncodeTextDeltaPayloadMatchesJSONMarshalForInvalidUTF8(t *testing.T) {
 	}
 }
 
+type fixedLatestVisibleMessageIDStore struct {
+	session.NoopStore
+	msgID int64
+}
+
+func (s *fixedLatestVisibleMessageIDStore) GetLatestVisibleMessageID(ctx context.Context, sessionID string) (int64, error) {
+	return s.msgID, nil
+}
+
+type blockingLatestVisibleMessageIDStore struct {
+	session.NoopStore
+	call func(context.Context)
+}
+
+func (s *blockingLatestVisibleMessageIDStore) GetLatestVisibleMessageID(ctx context.Context, sessionID string) (int64, error) {
+	if s.call != nil {
+		s.call(ctx)
+	}
+	<-ctx.Done()
+	return 0, ctx.Err()
+}
+
+func TestCompletedResponseIDForSessionUsesDurableIDWhenLookupSucceeds(t *testing.T) {
+	t.Parallel()
+
+	server := &serveServer{store: &fixedLatestVisibleMessageIDStore{msgID: 42}}
+	got := server.completedResponseIDForSessionWithin(context.Background(), "sess_test", "resp_ephemeral", time.Second)
+	want := durableResponseIDForMessageID(42)
+	if got != want {
+		t.Fatalf("completedResponseIDForSessionWithin() = %q, want %q", got, want)
+	}
+}
+
+func TestCompletedResponseIDForSessionFallsBackAfterBoundedLookup(t *testing.T) {
+	t.Parallel()
+
+	type ctxKey struct{}
+
+	var sawLiveCtx bool
+	var sawValue bool
+	server := &serveServer{store: &blockingLatestVisibleMessageIDStore{call: func(ctx context.Context) {
+		sawLiveCtx = ctx.Err() == nil
+		sawValue = ctx.Value(ctxKey{}) == "trace-123"
+	}}}
+
+	parentCtx, cancel := context.WithCancel(context.WithValue(context.Background(), ctxKey{}, "trace-123"))
+	cancel()
+
+	start := time.Now()
+	got := server.completedResponseIDForSessionWithin(parentCtx, "sess_test", "resp_ephemeral", 20*time.Millisecond)
+	elapsed := time.Since(start)
+
+	if got != "resp_ephemeral" {
+		t.Fatalf("completedResponseIDForSessionWithin() = %q, want fallback response id", got)
+	}
+	if !sawLiveCtx {
+		t.Fatal("durable lookup did not receive a live best-effort context")
+	}
+	if !sawValue {
+		t.Fatal("durable lookup did not preserve parent context values")
+	}
+	if elapsed < 10*time.Millisecond {
+		t.Fatalf("durable lookup returned too quickly (%v), want timeout-bounded best-effort lookup", elapsed)
+	}
+	if elapsed > 500*time.Millisecond {
+		t.Fatalf("durable lookup took too long (%v), want bounded fallback", elapsed)
+	}
+}
+
 func TestAppendResponseRunEventEmitsPhase(t *testing.T) {
 	run := newResponseRun("resp_phase", "sess_test", "", "mock", time.Now().Unix(), func() {})
 	server := &serveServer{}