Verify retry accounting and execution summary propagation
This commit is contained in:
@@ -1,5 +1,26 @@
|
|||||||
# Kit Upgrade Review Remediation
|
# Kit Upgrade Review Remediation
|
||||||
|
|
||||||
|
## Resolution
|
||||||
|
|
||||||
|
Completed 2026-09-26. No production changes were needed.
|
||||||
|
|
||||||
|
- `internal/framework/llm/promptkit_observation_test.go` now covers retries in
|
||||||
|
both initial and repair generations, aggregate versus legacy accounting,
|
||||||
|
unknown/known-zero/partial usage, retry-wait cancellation, exhausted execution,
|
||||||
|
preparation failure, and overlapping request isolation.
|
||||||
|
- `internal/framework/pipeline/debug_execution_test.go` exercises the real
|
||||||
|
scheduled client and debug wrapper together. It verifies returned and persisted
|
||||||
|
summaries for provider errors and cancellation without a body, split debug and
|
||||||
|
fallback content, terminal references, and explicit usage-presence JSON fields.
|
||||||
|
- Temporary regression probes confirmed that discarding failed responses in the
|
||||||
|
scheduler and collapsing usage presence into known zero cause the new tests
|
||||||
|
to fail. Those probes were reverted.
|
||||||
|
- Focused LLM/pipeline tests, their race-enabled suites, and
|
||||||
|
`GOWORK=off go test ./...` passed. Existing no-debug and redaction coverage was
|
||||||
|
retained. Runtime budgets, scheduler policy, and output contracts are unchanged.
|
||||||
|
|
||||||
|
The original finding below is retained as review context.
|
||||||
|
|
||||||
## Review Scope
|
## Review Scope
|
||||||
|
|
||||||
Reviewed implementation commits `31b48d2` through `7ac9d94` against
|
Reviewed implementation commits `31b48d2` through `7ac9d94` against
|
||||||
|
|||||||
245
internal/framework/llm/promptkit_observation_test.go
Normal file
245
internal/framework/llm/promptkit_observation_test.go
Normal file
@@ -0,0 +1,245 @@
|
|||||||
|
package llm
|
||||||
|
|
||||||
|
import (
|
||||||
|
"context"
|
||||||
|
"encoding/json"
|
||||||
|
"errors"
|
||||||
|
"io"
|
||||||
|
"net/http"
|
||||||
|
"net/http/httptest"
|
||||||
|
"strings"
|
||||||
|
"sync/atomic"
|
||||||
|
"testing"
|
||||||
|
"time"
|
||||||
|
|
||||||
|
"gitea.maximumdirect.net/eric/notarius/internal/framework/contracts"
|
||||||
|
"gitea.maximumdirect.net/eric/promptkit"
|
||||||
|
)
|
||||||
|
|
||||||
|
func newObservedHTTPClient(t *testing.T, handler http.HandlerFunc, retries int, transport http.RoundTripper) *PromptKitClient {
|
||||||
|
t.Helper()
|
||||||
|
server := httptest.NewServer(handler)
|
||||||
|
t.Cleanup(server.Close)
|
||||||
|
client, err := NewPromptKitClient(PromptKitClientConfig{
|
||||||
|
Assets: newTestPromptKitAssets(t), TransientRetries: &retries,
|
||||||
|
HTTPClient: &http.Client{Transport: transport},
|
||||||
|
EngineOptions: []promptkit.Option{promptkit.WithProfiles(promptkit.OpenAICompatibleProfile(promptkit.OpenAICompatibleProfileConfig{
|
||||||
|
ID: "default-profile", Endpoint: server.URL + "/v1", Model: "test-model",
|
||||||
|
}))},
|
||||||
|
})
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
return client
|
||||||
|
}
|
||||||
|
|
||||||
|
func observedRequest() contracts.StructuredCompletionRequest {
|
||||||
|
return contracts.StructuredCompletionRequest{
|
||||||
|
PromptID: "adapter.test", SessionID: "observation-test",
|
||||||
|
Inputs: contracts.LLMInputSet{"transcript": contracts.NewLLMInputMaterial("transcript", "application/json", []byte(`{}`), "", "")},
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestPromptKitObservationsComposeTransientRetriesAndRepair(t *testing.T) {
|
||||||
|
var calls atomic.Int32
|
||||||
|
client := newObservedHTTPClient(t, func(w http.ResponseWriter, r *http.Request) {
|
||||||
|
w.Header().Set("Content-Type", "application/json")
|
||||||
|
switch calls.Add(1) {
|
||||||
|
case 1, 3:
|
||||||
|
w.WriteHeader(http.StatusServiceUnavailable)
|
||||||
|
_, _ = io.WriteString(w, `{"error":{"message":"try again"}}`)
|
||||||
|
case 2:
|
||||||
|
_, _ = io.WriteString(w, `{"choices":[{"message":{"content":"{\"bad\":true}"}}],"usage":{"prompt_tokens":3,"completion_tokens":2,"total_tokens":5,"prompt_tokens_details":{"cached_tokens":0}}}`)
|
||||||
|
case 4:
|
||||||
|
_, _ = io.WriteString(w, `{"choices":[{"message":{"content":"{\"ok\":true}"}}],"usage":{"prompt_tokens":7,"completion_tokens":0,"total_tokens":7}}`)
|
||||||
|
default:
|
||||||
|
t.Error("unexpected extra HTTP attempt")
|
||||||
|
w.WriteHeader(http.StatusBadRequest)
|
||||||
|
}
|
||||||
|
}, 1, nil)
|
||||||
|
ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
|
||||||
|
defer cancel()
|
||||||
|
req := observedRequest()
|
||||||
|
repairs := 1
|
||||||
|
req.StructuredOutputRepairAttempts = &repairs
|
||||||
|
var out struct{ OK bool }
|
||||||
|
response, err := client.CompleteStructured(ctx, req, &out)
|
||||||
|
if err != nil || !out.OK {
|
||||||
|
t.Fatalf("completion = (%+v, %v), output = %+v", response, err, out)
|
||||||
|
}
|
||||||
|
summary := response.Execution
|
||||||
|
if calls.Load() != 4 || summary == nil || summary.Attempts != 4 || summary.TransientRetries != 2 || summary.RepairCalls != 1 || summary.Outcome != "success" || summary.RetriesExhausted {
|
||||||
|
t.Fatalf("calls = %d, execution = %+v", calls.Load(), summary)
|
||||||
|
}
|
||||||
|
wantUsage := contracts.LLMObservedTokenUsage{
|
||||||
|
PromptTokens: contracts.LLMObservedTokenCount{Value: 10, Known: true},
|
||||||
|
CompletionTokens: contracts.LLMObservedTokenCount{Value: 2, Known: true},
|
||||||
|
TotalTokens: contracts.LLMObservedTokenCount{Value: 12, Known: true},
|
||||||
|
CachedTokens: contracts.LLMObservedTokenCount{Known: true},
|
||||||
|
}
|
||||||
|
// Failed HTTP attempts contribute unknown usage, not known zero. Neither
|
||||||
|
// terminal totals nor legacy result totals should be added a second time.
|
||||||
|
if summary.Usage != wantUsage {
|
||||||
|
t.Fatalf("observed usage = %+v, want %+v", summary.Usage, wantUsage)
|
||||||
|
}
|
||||||
|
if response.PromptTokens != 10 || response.CompletionTokens != 2 || response.TotalTokens != 12 || response.RepairAttempts != 1 {
|
||||||
|
t.Fatalf("legacy result accounting changed: %+v", response)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
type closeSignalingBody struct {
|
||||||
|
io.ReadCloser
|
||||||
|
closed chan<- struct{}
|
||||||
|
}
|
||||||
|
|
||||||
|
func (b closeSignalingBody) Close() error {
|
||||||
|
err := b.ReadCloser.Close()
|
||||||
|
b.closed <- struct{}{}
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestPromptKitObservationSurvivesCanceledRetryWait(t *testing.T) {
|
||||||
|
consumed := make(chan struct{}, 2)
|
||||||
|
transport := http.DefaultTransport.(*http.Transport).Clone()
|
||||||
|
t.Cleanup(transport.CloseIdleConnections)
|
||||||
|
var calls atomic.Int32
|
||||||
|
client := newObservedHTTPClient(t, func(w http.ResponseWriter, r *http.Request) {
|
||||||
|
calls.Add(1)
|
||||||
|
w.Header().Set("Retry-After", "3600")
|
||||||
|
w.WriteHeader(http.StatusServiceUnavailable)
|
||||||
|
_, _ = io.WriteString(w, `{"error":{"message":"busy"}}`)
|
||||||
|
}, 1, roundTripFunc(func(req *http.Request) (*http.Response, error) {
|
||||||
|
response, err := transport.RoundTrip(req)
|
||||||
|
if err == nil {
|
||||||
|
response.Body = closeSignalingBody{ReadCloser: response.Body, closed: consumed}
|
||||||
|
}
|
||||||
|
return response, err
|
||||||
|
}))
|
||||||
|
ctx, cancel := context.WithCancel(context.Background())
|
||||||
|
defer cancel()
|
||||||
|
type result struct {
|
||||||
|
response contracts.StructuredCompletionResponse
|
||||||
|
err error
|
||||||
|
}
|
||||||
|
done := make(chan result, 1)
|
||||||
|
go func() {
|
||||||
|
var out map[string]any
|
||||||
|
response, err := client.CompleteStructured(ctx, observedRequest(), &out)
|
||||||
|
done <- result{response, err}
|
||||||
|
}()
|
||||||
|
// The failure has been consumed, so cancellation must interrupt the pending
|
||||||
|
// retry wait rather than depend on another provider response or its timeout.
|
||||||
|
select {
|
||||||
|
case <-consumed:
|
||||||
|
case <-time.After(5 * time.Second):
|
||||||
|
t.Fatal("provider response was not consumed")
|
||||||
|
}
|
||||||
|
cancel()
|
||||||
|
select {
|
||||||
|
case got := <-done:
|
||||||
|
if !errors.Is(got.err, context.Canceled) || got.response.Execution == nil {
|
||||||
|
t.Fatalf("canceled completion = (%+v, %v)", got.response, got.err)
|
||||||
|
}
|
||||||
|
summary := got.response.Execution
|
||||||
|
if summary.Outcome != "error" || summary.FailureKind != "canceled" || summary.Attempts != 1 || summary.RetriesExhausted || calls.Load() != 1 {
|
||||||
|
t.Fatalf("canceled execution = %+v, HTTP calls = %d", summary, calls.Load())
|
||||||
|
}
|
||||||
|
case <-time.After(5 * time.Second):
|
||||||
|
t.Fatal("cancellation did not interrupt retry waiting")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestPromptKitObservationsSeparatePreparationAndExecutionFailure(t *testing.T) {
|
||||||
|
var calls atomic.Int32
|
||||||
|
client := newObservedHTTPClient(t, func(w http.ResponseWriter, r *http.Request) {
|
||||||
|
calls.Add(1)
|
||||||
|
w.WriteHeader(http.StatusServiceUnavailable)
|
||||||
|
_, _ = io.WriteString(w, `{"error":{"message":"busy"}}`)
|
||||||
|
}, 1, nil)
|
||||||
|
ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
|
||||||
|
defer cancel()
|
||||||
|
req := observedRequest()
|
||||||
|
req.PromptID = "missing-prompt"
|
||||||
|
var out map[string]any
|
||||||
|
response, err := client.CompleteStructured(ctx, req, &out)
|
||||||
|
if err == nil || response.Execution != nil || calls.Load() != 0 {
|
||||||
|
t.Fatalf("preparation failure = (%+v, %v), calls = %d", response, err, calls.Load())
|
||||||
|
}
|
||||||
|
response, err = client.CompleteStructured(ctx, observedRequest(), &out)
|
||||||
|
var generation *contracts.LLMGenerationError
|
||||||
|
if !errors.As(err, &generation) || response.Execution == nil {
|
||||||
|
t.Fatalf("execution failure = (%+v, %v)", response, err)
|
||||||
|
}
|
||||||
|
summary := response.Execution
|
||||||
|
if summary.Outcome != "error" || summary.FailureKind != "provider_status" || summary.StatusCode != http.StatusServiceUnavailable || !summary.RetriesExhausted || summary.Attempts != 2 || summary.TransientRetries != 1 || calls.Load() != 2 {
|
||||||
|
t.Fatalf("failed execution = %+v, calls = %d", summary, calls.Load())
|
||||||
|
}
|
||||||
|
if summary.Usage != (contracts.LLMObservedTokenUsage{}) {
|
||||||
|
t.Fatalf("failed attempts should have unknown usage: %+v", summary.Usage)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestPromptKitObservationsAreIsolatedAcrossConcurrentCalls(t *testing.T) {
|
||||||
|
arrived := make(chan struct{}, 2)
|
||||||
|
release := make(chan struct{})
|
||||||
|
client := newObservedHTTPClient(t, func(w http.ResponseWriter, r *http.Request) {
|
||||||
|
body, err := io.ReadAll(r.Body)
|
||||||
|
if err != nil {
|
||||||
|
t.Error(err)
|
||||||
|
return
|
||||||
|
}
|
||||||
|
tokens := 11
|
||||||
|
if strings.Contains(string(body), "second-call") {
|
||||||
|
tokens = 22
|
||||||
|
}
|
||||||
|
arrived <- struct{}{}
|
||||||
|
select {
|
||||||
|
case <-release:
|
||||||
|
case <-r.Context().Done():
|
||||||
|
return
|
||||||
|
}
|
||||||
|
_ = json.NewEncoder(w).Encode(map[string]any{
|
||||||
|
"choices": []any{map[string]any{"message": map[string]string{"content": `{"ok":true}`}}},
|
||||||
|
"usage": map[string]int{"prompt_tokens": tokens, "completion_tokens": 0, "total_tokens": tokens},
|
||||||
|
})
|
||||||
|
}, 0, nil)
|
||||||
|
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
|
||||||
|
defer cancel()
|
||||||
|
type result struct {
|
||||||
|
want int
|
||||||
|
response contracts.StructuredCompletionResponse
|
||||||
|
err error
|
||||||
|
}
|
||||||
|
done := make(chan result, 2)
|
||||||
|
for name, tokens := range map[string]int{"first-call": 11, "second-call": 22} {
|
||||||
|
go func() {
|
||||||
|
req := observedRequest()
|
||||||
|
req.Vars = map[string]any{"custom": name}
|
||||||
|
var out map[string]any
|
||||||
|
response, err := client.CompleteStructured(ctx, req, &out)
|
||||||
|
done <- result{tokens, response, err}
|
||||||
|
}()
|
||||||
|
}
|
||||||
|
for i := 0; i < 2; i++ {
|
||||||
|
select {
|
||||||
|
case <-arrived:
|
||||||
|
case <-ctx.Done():
|
||||||
|
t.Fatal("requests did not overlap")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
close(release)
|
||||||
|
for i := 0; i < 2; i++ {
|
||||||
|
select {
|
||||||
|
case got := <-done:
|
||||||
|
summary := got.response.Execution
|
||||||
|
if got.err != nil || summary == nil {
|
||||||
|
t.Fatalf("concurrent completion = (%+v, %v)", got.response, got.err)
|
||||||
|
}
|
||||||
|
if summary.Attempts != 1 || summary.Usage.PromptTokens != (contracts.LLMObservedTokenCount{Value: got.want, Known: true, Complete: true}) || summary.Usage.CompletionTokens != (contracts.LLMObservedTokenCount{Known: true, Complete: true}) || summary.Usage.CachedTokens != (contracts.LLMObservedTokenCount{}) {
|
||||||
|
t.Fatalf("execution for %d tokens = %+v", got.want, summary)
|
||||||
|
}
|
||||||
|
case <-ctx.Done():
|
||||||
|
t.Fatal("concurrent completions did not return")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
122
internal/framework/pipeline/debug_execution_test.go
Normal file
122
internal/framework/pipeline/debug_execution_test.go
Normal file
@@ -0,0 +1,122 @@
|
|||||||
|
package pipeline
|
||||||
|
|
||||||
|
import (
|
||||||
|
"bytes"
|
||||||
|
"context"
|
||||||
|
"encoding/json"
|
||||||
|
"errors"
|
||||||
|
"testing"
|
||||||
|
|
||||||
|
"gitea.maximumdirect.net/eric/notarius/internal/framework/contracts"
|
||||||
|
"gitea.maximumdirect.net/eric/notarius/internal/framework/llm"
|
||||||
|
)
|
||||||
|
|
||||||
|
type executionDebugClient struct {
|
||||||
|
response contracts.StructuredCompletionResponse
|
||||||
|
err error
|
||||||
|
}
|
||||||
|
|
||||||
|
func (c executionDebugClient) CompleteStructured(context.Context, contracts.StructuredCompletionRequest, any) (contracts.StructuredCompletionResponse, error) {
|
||||||
|
return c.response, c.err
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestScheduledDebugClientPreservesExecutionThroughTerminalRecords(t *testing.T) {
|
||||||
|
providerErr := contracts.NewLLMGenerationError(503, "provider unavailable")
|
||||||
|
const content = `{"ok":true}`
|
||||||
|
for _, test := range []struct {
|
||||||
|
name string
|
||||||
|
err error
|
||||||
|
content string
|
||||||
|
debugResponse bool
|
||||||
|
}{
|
||||||
|
{name: "provider failure without body", err: providerErr},
|
||||||
|
{name: "canceled without body", err: context.Canceled},
|
||||||
|
{name: "split debug response", content: content, debugResponse: true},
|
||||||
|
{name: "split fallback response", content: content},
|
||||||
|
} {
|
||||||
|
t.Run(test.name, func(t *testing.T) {
|
||||||
|
want := contracts.LLMExecutionSummary{
|
||||||
|
Attempts: 2, TransientRetries: 1, Outcome: "success",
|
||||||
|
CapacityWaitMS: 1, GenerationMS: 3, RetryDelayMS: 5, DurationMS: 9,
|
||||||
|
Usage: contracts.LLMObservedTokenUsage{
|
||||||
|
PromptTokens: contracts.LLMObservedTokenCount{Value: 7, Known: true},
|
||||||
|
CompletionTokens: contracts.LLMObservedTokenCount{Known: true},
|
||||||
|
},
|
||||||
|
}
|
||||||
|
if test.err != nil {
|
||||||
|
want.Outcome = "error"
|
||||||
|
want.FailureKind = "provider_status"
|
||||||
|
want.StatusCode = 503
|
||||||
|
want.RetriesExhausted = true
|
||||||
|
if errors.Is(test.err, context.Canceled) {
|
||||||
|
want.FailureKind = "canceled"
|
||||||
|
want.StatusCode = 0
|
||||||
|
want.RetriesExhausted = false
|
||||||
|
}
|
||||||
|
}
|
||||||
|
execution := want
|
||||||
|
response := contracts.StructuredCompletionResponse{Execution: &execution}
|
||||||
|
if test.debugResponse {
|
||||||
|
response.Debug = &contracts.LLMDebugMaterial{Response: &contracts.LLMDebugResponse{Content: test.content}}
|
||||||
|
} else {
|
||||||
|
response.Content = json.RawMessage(test.content)
|
||||||
|
}
|
||||||
|
scheduler, err := llm.NewScheduler(1)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
recorder := newCapturedDebugRecorder()
|
||||||
|
client := WithDebugLLMRecording(llm.NewScheduledClient(executionDebugClient{response, test.err}, scheduler), recorder)
|
||||||
|
ctx, scope := withDebugLLMScope(context.Background(), "attempt")
|
||||||
|
got, err := client.CompleteStructured(ctx, contracts.StructuredCompletionRequest{StageName: "extract"}, &struct{}{})
|
||||||
|
if !errors.Is(err, test.err) || got.Execution == nil || *got.Execution != want {
|
||||||
|
t.Fatalf("scheduled debug completion = (%+v, %v), want summary and original error", got, err)
|
||||||
|
}
|
||||||
|
terminal := newAttemptTerminalRecorder(recorder, "attempt/terminal", "extract", scope, debugTimedEnvelope{Stage: "extract"})
|
||||||
|
if err := terminal.record(nil, test.err); !errors.Is(err, test.err) {
|
||||||
|
t.Fatalf("terminal recording error = %v", err)
|
||||||
|
}
|
||||||
|
envelope := recorder.envelope(t, "attempt/response-0001.json")
|
||||||
|
encoded, err := json.Marshal(envelope.Payload)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
var artifact debugLLMResponseArtifact
|
||||||
|
if err := json.Unmarshal(encoded, &artifact); err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
if artifact.Execution == nil || *artifact.Execution != want {
|
||||||
|
t.Fatalf("persisted response lost execution summary: %+v", artifact)
|
||||||
|
}
|
||||||
|
if !test.debugResponse && (artifact.Fallback == nil || artifact.Fallback.Execution == nil || *artifact.Fallback.Execution != want) {
|
||||||
|
t.Fatalf("persisted fallback lost execution summary: %+v", artifact.Fallback)
|
||||||
|
}
|
||||||
|
if test.content == "" {
|
||||||
|
if artifact.ContentPath != "" {
|
||||||
|
t.Fatalf("bodyless failure has content path %q", artifact.ContentPath)
|
||||||
|
}
|
||||||
|
} else {
|
||||||
|
var compact bytes.Buffer
|
||||||
|
if artifact.ContentPath == "" || json.Compact(&compact, recorder.bytes[artifact.ContentPath]) != nil || compact.String() != test.content {
|
||||||
|
t.Fatalf("split content not preserved at %q", artifact.ContentPath)
|
||||||
|
}
|
||||||
|
if (artifact.Fallback != nil && artifact.Fallback.Content != "") || (artifact.Response != nil && artifact.Response.Content != "") {
|
||||||
|
t.Fatal("split content remains in response metadata")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
refs := recorder.envelope(t, "attempt/terminal.json").LLMCalls
|
||||||
|
if len(refs) != 1 || refs[0].Execution == nil || *refs[0].Execution != want || refs[0].Error != (test.err != nil) {
|
||||||
|
t.Fatalf("terminal references = %+v", refs)
|
||||||
|
}
|
||||||
|
// Decode the persisted JSON, not just Go defaults: omitted false and
|
||||||
|
// zero fields would erase the presence-aware usage contract.
|
||||||
|
usage := envelope.Payload.(map[string]any)["execution"].(map[string]any)["usage"].(map[string]any)
|
||||||
|
unknown := usage["cached_tokens"].(map[string]any)
|
||||||
|
for _, key := range []string{"value", "known", "complete"} {
|
||||||
|
if _, ok := unknown[key]; !ok {
|
||||||
|
t.Fatalf("unknown token usage omits %q: %+v", key, unknown)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
})
|
||||||
|
}
|
||||||
|
}
|
||||||
Reference in New Issue
Block a user