Share generation attempt accounting

internal/generation/coordinator.go

@@ -130,36 +130,7 @@ func (c *Coordinator) generate(
 		if result.Dispatched {
 			attempts++
 		}
-		if captureEvidence {
-			attempt := AttemptEvidence{
-				StartedAt:           result.StartedAt.UTC(),
-				Dispatched:          result.Dispatched,
-				ResponseReceived:    result.Outcome.Evidence.ResponseReceived,
-				HTTP:                cloneHTTPEvidence(result.Outcome.Evidence),
-				CapacityWait:        result.CapacityWait,
-				RequestDuration:     result.RequestDuration,
-				Failure:             ClassifyFailure(result.Outcome.Err),
-				ResponseDecoded:     result.Outcome.ResponseDecoded,
-				GenerationCompleted: result.Outcome.GenerationCompleted,
-			}
-			if result.Outcome.StructuredOutputRequested != nil {
-				requested := *result.Outcome.StructuredOutputRequested
-				attempt.StructuredOutputRequested = &requested
-			}
-			if result.Dispatched {
-				usage.Add(result.Outcome.Usage, result.Outcome.UsagePresence)
-				var attemptUsage observation.UsageAccumulator
-				attemptUsage.Add(result.Outcome.Usage, result.Outcome.UsagePresence)
-				attempt.Usage = attemptUsage.Observed()
-			}
-			evidence.Attempts = append(evidence.Attempts, attempt)
-		}
-		if recorder != nil {
-			recorder.Attempt(observation.AttemptData{Generation: generation, Attempt: ordinal,
-				Dispatched: result.Dispatched, CapacityWait: result.CapacityWait, GenerationDuration: result.RequestDuration,
-				Outcome: attemptOutcome(result.Outcome.Err), Failure: ClassifyFailure(result.Outcome.Err),
-				StatusCode: result.Outcome.StatusCode, Usage: result.Outcome.Usage, UsagePresence: result.Outcome.UsagePresence})
-		}
+		recordAttempt(result, false, captureEvidence, &evidence, &usage, recorder, generation, ordinal)
 		if result.Outcome.Err == nil {
 			if err := ctx.Err(); err != nil {
 				if captureEvidence {
@@ -206,104 +177,57 @@ func (c *Coordinator) generate(
 	}
 }
 
-func (c *Coordinator) generateInjected(
-	ctx context.Context,
-	req domain.GenerateRequest,
-	captureEvidence bool,
-) (*domain.GenerateResponse, error, GenerationEvidence) {
+func (c *Coordinator) generateInjected(ctx context.Context, req domain.GenerateRequest, captureEvidence bool) (*domain.GenerateResponse, error, GenerationEvidence) {
 	evidence := GenerationEvidence{}
 	var usage observation.UsageAccumulator
-	finish := func(response *domain.GenerateResponse, err error) (*domain.GenerateResponse, error, GenerationEvidence) {
-		if captureEvidence {
-			evidence.Usage = usage.Observed()
-		}
-		return response, err, evidence
-	}
-	recorder := observation.FromContext(ctx)
 	generation := req.Generation
 	if generation == 0 {
 		generation = 1
 	}
-	start := c.now()
-	release, err := c.manager.Acquire(ctx, req.Target.BackendID)
-	capacityWait := c.capacityWait(req.Target.BackendID, start)
-	if err != nil {
+	result := c.executeInjected(ctx, req)
+	recordAttempt(result, true, captureEvidence, &evidence, &usage, observation.FromContext(ctx), generation, 1)
+	if captureEvidence {
+		evidence.Usage = usage.Observed()
+	}
+	if result.Outcome.Err != nil {
+		return nil, result.Outcome.Err, evidence
+	}
+	if err := ctx.Err(); err != nil {
 		if captureEvidence {
-			evidence.Attempts = append(evidence.Attempts, AttemptEvidence{
-				StartedAt:    start.UTC(),
-				CapacityWait: capacityWait,
-				Failure:      ClassifyFailure(err),
-			})
+			return result.Outcome.Response, err, evidence
 		}
-		if recorder != nil {
-			recorder.Attempt(observation.AttemptData{Generation: generation, Attempt: 1, CapacityWait: capacityWait, Outcome: domain.ExecutionOutcomeError, Failure: ClassifyFailure(err)})
-		}
-		return finish(nil, err)
+		return nil, err, evidence
+	}
+	return result.Outcome.Response, nil, evidence
+}
+
+func (c *Coordinator) executeInjected(ctx context.Context, req domain.GenerateRequest) attemptResult {
+	result := attemptResult{StartedAt: c.now()}
+	release, err := c.manager.Acquire(ctx, req.Target.BackendID)
+	result.CapacityWait = c.capacityWait(req.Target.BackendID, result.StartedAt)
+	if err != nil {
+		result.Outcome.Err = err
+		return result
 	}
 	defer release()
 	if err := ctx.Err(); err != nil {
-		release()
-		if captureEvidence {
-			evidence.Attempts = append(evidence.Attempts, AttemptEvidence{
-				StartedAt:    start.UTC(),
-				CapacityWait: capacityWait,
-				Failure:      ClassifyFailure(err),
-			})
-		}
-		if recorder != nil {
-			recorder.Attempt(observation.AttemptData{Generation: generation, Attempt: 1, CapacityWait: capacityWait, Outcome: domain.ExecutionOutcomeError, Failure: ClassifyFailure(err)})
-		}
-		return finish(nil, err)
+		result.Outcome.Err = err
+		return result
 	}
-	dispatchStart := c.now()
-	response, err := c.injected.Generate(ctx, req)
-	duration := elapsed(dispatchStart, c.now())
-	if err == nil && response == nil {
-		err = llm.ErrMalformedResponse
+	start := c.now()
+	result.Dispatched = true
+	requested := req.StructuredOutput != nil
+	result.Outcome.StructuredOutputRequested = &requested
+	result.Outcome.Response, result.Outcome.Err = c.injected.Generate(ctx, req)
+	result.RequestDuration = elapsed(start, c.now())
+	if result.Outcome.Err == nil && result.Outcome.Response == nil {
+		result.Outcome.Err = llm.ErrMalformedResponse
 	}
-	if captureEvidence {
-		attempt := AttemptEvidence{
-			StartedAt:           start.UTC(),
-			Dispatched:          true,
-			ResponseReceived:    response != nil,
-			RequestDuration:     duration,
-			Failure:             ClassifyFailure(err),
-			GenerationCompleted: err == nil && response != nil,
-		}
-		requested := req.StructuredOutput != nil
-		attempt.StructuredOutputRequested = &requested
-		var reportedUsage domain.TokenUsage
-		var reportedPresence *domain.TokenUsagePresence
-		if err == nil && response != nil {
-			reportedUsage = response.Usage
-			reportedPresence = response.UsagePresence
-		}
-		usage.Add(reportedUsage, reportedPresence)
-		var attemptUsage observation.UsageAccumulator
-		attemptUsage.Add(reportedUsage, reportedPresence)
-		attempt.Usage = attemptUsage.Observed()
-		evidence.Attempts = append(evidence.Attempts, attempt)
+	if result.Outcome.Err == nil {
+		result.Outcome.Usage = result.Outcome.Response.Usage
+		result.Outcome.UsagePresence = result.Outcome.Response.UsagePresence
 	}
-	release()
-	if recorder != nil {
-		data := observation.AttemptData{Generation: generation, Attempt: 1, Dispatched: true, CapacityWait: capacityWait, GenerationDuration: duration,
-			Outcome: attemptOutcome(err), Failure: ClassifyFailure(err)}
-		if err == nil {
-			data.Usage = response.Usage
-			data.UsagePresence = response.UsagePresence
-		}
-		recorder.Attempt(data)
-	}
-	if err != nil {
-		return finish(nil, err)
-	}
-	if err := ctx.Err(); err != nil {
-		if captureEvidence {
-			return finish(response, err)
-		}
-		return finish(nil, err)
-	}
-	return finish(response, nil)
+	return result
 }
 
 type attemptResult struct {
@@ -314,6 +238,49 @@ type attemptResult struct {
 	StartedAt       time.Time
 }
 
+// recordAttempt runs after the attempt has released its permit, including when
+// the recorder calls application code. Injected clients expose no HTTP facts.
+func recordAttempt(result attemptResult, injected, capture bool, evidence *GenerationEvidence, usage *observation.UsageAccumulator, recorder *observation.Recorder, generation, ordinal uint64) {
+	outcome := result.Outcome
+	failure := ClassifyFailure(outcome.Err)
+	if capture {
+		attempt := AttemptEvidence{
+			StartedAt: result.StartedAt.UTC(), Dispatched: result.Dispatched,
+			CapacityWait: result.CapacityWait, RequestDuration: result.RequestDuration,
+			Failure: failure,
+		}
+		if injected {
+			attempt.ResponseReceived = outcome.Response != nil
+			attempt.GenerationCompleted = outcome.Err == nil && outcome.Response != nil
+			if result.Dispatched {
+				attempt.CapacityWait = 0
+			}
+		} else {
+			attempt.ResponseReceived = outcome.Evidence.ResponseReceived
+			attempt.HTTP = cloneHTTPEvidence(outcome.Evidence)
+			attempt.ResponseDecoded = outcome.ResponseDecoded
+			attempt.GenerationCompleted = outcome.GenerationCompleted
+		}
+		if outcome.StructuredOutputRequested != nil {
+			requested := *outcome.StructuredOutputRequested
+			attempt.StructuredOutputRequested = &requested
+		}
+		if result.Dispatched {
+			usage.Add(outcome.Usage, outcome.UsagePresence)
+			var single observation.UsageAccumulator
+			single.Add(outcome.Usage, outcome.UsagePresence)
+			attempt.Usage = single.Observed()
+		}
+		evidence.Attempts = append(evidence.Attempts, attempt)
+	}
+	if recorder != nil {
+		recorder.Attempt(observation.AttemptData{Generation: generation, Attempt: ordinal,
+			Dispatched: result.Dispatched, CapacityWait: result.CapacityWait, GenerationDuration: result.RequestDuration,
+			Outcome: attemptOutcome(outcome.Err), Failure: failure, StatusCode: outcome.StatusCode,
+			Usage: outcome.Usage, UsagePresence: outcome.UsagePresence})
+	}
+}
+
 func (c *Coordinator) execute(ctx context.Context, backendID string, prepared *llm.PreparedGeneration, captureEvidence bool) attemptResult {
 	result := attemptResult{StartedAt: c.now()}
 	release, err := c.manager.Acquire(ctx, backendID)

internal/generation/observation_test.go

@@ -144,6 +144,32 @@ func (c fixedClient) Generate(context.Context, domain.GenerateRequest) (*domain.
 	return c.response, c.err
 }
 
+func TestInjectedObserverRunsAfterPermitRelease(t *testing.T) {
+	configured := testCoordinator(t, func(*http.Request) (*http.Response, error) { t.Fatal("unexpected HTTP request"); return nil, nil })
+	coordinator := NewInjected(configured.manager, fixedClient{response: &domain.GenerateResponse{}})
+	called := false
+	recorder, err := observation.New(func(event domain.ExecutionEvent) {
+		if event.Kind != domain.ExecutionAttempt {
+			return
+		}
+		called = true
+		ctx, cancel := context.WithTimeout(context.Background(), time.Second)
+		defer cancel()
+		release, err := coordinator.manager.Acquire(ctx, "limited")
+		if err != nil {
+			t.Errorf("observer could not acquire permit: %v", err)
+			return
+		}
+		release()
+	})
+	if err != nil {
+		t.Fatal(err)
+	}
+	if _, err := coordinator.Generate(observation.WithRecorder(context.Background(), recorder), testRequest(0)); err != nil || !called {
+		t.Fatalf("generation error = %v, observer called = %v", err, called)
+	}
+}
+
 func TestObservedInjectedUsageAndRepairGeneration(t *testing.T) {
 	manager, err := capacity.NewManager(nil)
 	if err != nil {