diff --git a/cmd/serve_response_runs.go b/cmd/serve_response_runs.go index b6be12c4..a9a8ce43 100644 --- a/cmd/serve_response_runs.go +++ b/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) } diff --git a/cmd/serve_response_runs_test.go b/cmd/serve_response_runs_test.go index d9a1eaf5..c94f8576 100644 --- a/cmd/serve_response_runs_test.go +++ b/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{}