Fix prediction response timestamp population and webhook finalization (#219)

* Fix prediction response timestamp population and webhook finalization

This PR addresses critical issues with timestamp population in prediction responses and consolidates webhook finalization logic to ensure consistent behavior across sync and async predictions.

Core Infrastructure Changes

Time Format Standardization
- Add TimeFormat constant to config package using microsecond precision with explicit timezone offset
- Ensures consistent timestamp formatting across all prediction responses and webhook payloads

PredictionResponse Structure Overhaul
- Remove deprecated WebhookURL field from PredictionResponse struct
- Add dedicated CreatedAt, StartedAt, CompletedAt timestamp fields with proper JSON tags
- Change Metrics field from any to map[string]any for type safety
- Add finalizeResponse() method to set CompletedAt timestamp and calculate predict_time metric
- Add populateFromRequest() method to copy core fields (ID, Input, CreatedAt, StartedAt) from request

Prediction Lifecycle Management

Manager API Restructuring
- Rename Manager.Predict() to Manager.PredictSync() for clarity
- Update Manager.PredictAsync() to return initial populated response instead of void
- Refactor internal predict() method to return both channel and initial response
- Add async parameter to assignReqToRunner() to control setup wait behavior

Runner Prediction Flow
- Update runner.predict() signature to return initial response populated with request fields
- Fix timestamp population in predict() method by setting CreatedAt/StartedAt on pending request
- Add proper populateFromRequest() calls in handleResponseWebhooksAndCompletion after response overwrites
- Ensure terminal responses call finalizeResponse() before sending through safeSend()

Webhook and Response Handling
- Consolidate terminal webhook logic into unified sendTerminalWebhook() method
- Fix race conditions by ensuring response finalization before webhook sending
- Update subprocess monitoring to use populateFromRequest() for failed prediction responses
- Ensure all webhook payloads include complete timestamp information

Test Infrastructure Improvements

Response Validation Framework
- Add comprehensive unit tests for finalizeResponse() covering edge cases and error conditions
- Add unit tests for populateFromRequest() method with field overwrite scenarios
- Update all integration tests to validate CreatedAt, StartedAt, CompletedAt fields
- Add ValidateTerminalWebhookResponse helper function for consistent terminal response validation

Test Harness Enhancements
- Introduce testHarnessResponse type with string Logs field for easier test assertions
- Update webhook receiver to use testHarnessResponse for simplified log comparisons
- Retrofit existing tests to use new validation patterns without changing test logic

Bug Fixes

Timestamp Population Issues
- Fix missing timestamps in sync prediction responses by ensuring populateFromRequest() calls
- Resolve timestamp loss during response overwrites in handleResponseWebhooksAndCompletion
- Ensure webhook payloads contain complete timestamp information for terminal events

Metrics and Finalization
- Fix missing predict_time metric in terminal responses by calling finalizeResponse()
- Ensure Metrics map initialization before metric calculation
- Handle time parsing errors gracefully in finalizeResponse() method

Concurrency and Error Handling
- Fix potential race conditions in webhook sending by proper response finalization order
- Ensure cleanup of failed predictions includes populated response fields
- Update error responses in runner stop scenarios to include request-derived fields

Backward Compatibility

- Maintain existing API signatures where possible
- Preserve webhook payload structure while adding timestamp fields
- Keep existing test behavior while enhancing validation capabilities

* Refactor handleResponseWebhooksAndCompletion to eliminate duplicate response handling

- Consolidate response update logic into single operation to avoid duplicate overwrites
- Extract webhook sending logic into dedicated sendStatusWebhook function
- Extract terminal completion handling into handleTerminalCompletion function
- Ensure proper lock context around populateFromRequest calls accessing pending.request
- Maintain race condition protection by using local response copy for safeSend
- Remove TODO comment about function being "a mess" - now clean and maintainable

This eliminates the scattered response object handling and reduces lock contention
while preserving all existing functionality and thread safety guarantees.

* nit fixes

Author: tempusfrangit

internal/config/config.go

@@ -5,6 +5,10 @@ import (
 	"time"
 )
 
+const (
+	TimeFormat = "2006-01-02T15:04:05.999999-07:00"
+)
+
 // Config holds all configuration for the cog runtime service
 type Config struct {
 	// Server configuration

internal/runner/manager.go

@@ -1,10 +1,8 @@
 package runner
 
 import (
-	"bytes"
 	"context"
 	_ "embed"
-	"encoding/json"
 	"errors"
 	"fmt"
 	"io/fs"
@@ -152,9 +150,9 @@ func (m *Manager) releaseSlot() {
 	}
 }
 
-// Predict executes a sync prediction request - blocks until complete
-func (m *Manager) Predict(req PredictionRequest) (*PredictionResponse, error) {
-	respChan, err := m.predict(m.ctx, req)
+// PredictSync executes a sync prediction request - blocks until complete
+func (m *Manager) PredictSync(req PredictionRequest) (*PredictionResponse, error) {
+	respChan, _, err := m.predict(m.ctx, req, false)
 	if err != nil {
 		return nil, err
 	}
@@ -165,41 +163,11 @@ func (m *Manager) Predict(req PredictionRequest) (*PredictionResponse, error) {
 }
 
 // PredictAsync executes an async prediction request - returns immediately, sends webhook when complete
-func (m *Manager) PredictAsync(ctx context.Context, req PredictionRequest) error {
+func (m *Manager) PredictAsync(ctx context.Context, req PredictionRequest) (*PredictionResponse, error) {
 	log := m.logger.Sugar()
-	if err := m.claimSlot(); err != nil {
-		return err
-	}
-
-	deadlineCtx, deadlineCancel := context.WithTimeout(ctx, 10*time.Second)
-	defer deadlineCancel()
-
-	runner, err := m.assignReqToRunner(deadlineCtx, req)
+	respChan, initialResponse, err := m.predict(ctx, req, true)
 	if err != nil {
-		log.Tracew("failed to get runner for async request", "error", err)
-		m.releaseSlot()
-		return err
-	}
-
-	switch runner.status {
-	case StatusReady:
-		// Status ready is always valid for new predictions
-		break
-	case StatusStarting:
-		if !m.cfg.UseProcedureMode {
-			m.releaseSlot()
-			return fmt.Errorf("%w: %s", ErrInvalidRunnerStatus, runner.status)
-		}
-	default:
-		m.releaseSlot()
-		return fmt.Errorf("%w: %s", ErrInvalidRunnerStatus, runner.status)
-	}
-
-	respChan, err := runner.predict(req)
-	if err != nil {
-		log.Tracew("failed to predict", "error", err)
-		m.releaseSlot()
-		return err
+		return nil, err
 	}
 
 	// Release slot when prediction completes in background
@@ -209,33 +177,33 @@ func (m *Manager) PredictAsync(ctx context.Context, req PredictionRequest) error
 		log.Tracew("async prediction completed", "prediction_id", req.ID)
 	}()
 
-	return nil
+	return initialResponse, nil
 }
 
 // predict is the internal implementation shared by both sync and async predictions
-func (m *Manager) predict(ctx context.Context, req PredictionRequest) (chan PredictionResponse, error) {
+func (m *Manager) predict(ctx context.Context, req PredictionRequest, async bool) (chan PredictionResponse, *PredictionResponse, error) {
 	if err := m.claimSlot(); err != nil {
-		return nil, err
+		return nil, nil, err
 	}
 
 	deadlineCtx, deadlineCancel := context.WithTimeout(ctx, 10*time.Second)
 	defer deadlineCancel()
 
-	runner, err := m.assignReqToRunnerWait(deadlineCtx, req)
+	runner, err := m.assignReqToRunner(deadlineCtx, req, async)
 	if err != nil {
 		m.releaseSlot()
-		return nil, err
+		return nil, nil, err
 	}
 
 	if !m.cfg.UseProcedureMode && runner.status != StatusReady {
 		m.releaseSlot()
-		return nil, fmt.Errorf("runner not ready: %s", runner.status)
+		return nil, nil, fmt.Errorf("runner not ready: %s", runner.status)
 	}
 
-	respChan, err := runner.predict(req)
+	respChan, initialResponse, err := runner.predict(req)
 	if err != nil {
 		m.releaseSlot()
-		return nil, err
+		return nil, nil, err
 	}
 
 	// Wrap the channel to release slot when prediction completes
@@ -247,39 +215,7 @@ func (m *Manager) predict(ctx context.Context, req PredictionRequest) (chan Pred
 		close(wrappedChan)
 	}()
 
-	return wrappedChan, nil
-}
-
-// sendTerminalWebhook sends a terminal webhook synchronously for a completed prediction
-func (m *Manager) sendTerminalWebhook(req PredictionRequest, resp PredictionResponse) error {
-	log := m.logger.Sugar()
-
-	// Send synchronously using SendConditional to respect filters
-	// Send the actual PredictionResponse object, not a custom map
-
-	body, err := json.Marshal(resp)
-	if err != nil {
-		log.Errorw("failed to marshal prediction response", "error", err)
-		return fmt.Errorf("failed to marshal prediction response: %w", err)
-	}
-
-	if err := m.webhookSender.SendConditional(req.Webhook, bytes.NewReader(body), webhook.EventCompleted, req.WebhookEventsFilter, nil); err != nil {
-		log.Errorw("failed to send terminal webhook", "prediction_id", resp.ID, "webhook_url", req.Webhook, "error", err)
-		return fmt.Errorf("failed to send terminal webhook: %w", err)
-	}
-	log.Infow("sent terminal webhook", "prediction_id", resp.ID, "webhook_url", req.Webhook, "status", resp.Status)
-	return nil
-}
-
-func (m *Manager) assignReqToRunnerWait(ctx context.Context, req PredictionRequest) (*Runner, error) {
-	runner, err := m.assignReqToRunner(ctx, req)
-	if err != nil {
-		return nil, fmt.Errorf("failed to assign request to runner: %w", err)
-	}
-	if waitForRunnerSetup(ctx, runner) != nil {
-		return nil, err
-	}
-	return runner, nil
+	return wrappedChan, initialResponse, nil
 }
 
 // createDefaultRunner creates the default runner for non-procedure mode
@@ -384,6 +320,7 @@ func (m *Manager) createDefaultRunner(ctx context.Context) (*Runner, error) {
 
 // allocatePrediction reserves a slot in the runner for the prediction
 func (m *Manager) allocatePrediction(runner *Runner, req PredictionRequest) { //nolint:contextcheck // we do not use this context for the prediction see note below
+	log := m.logger.Sugar()
 	runner.mu.Lock()
 	defer runner.mu.Unlock()
 
@@ -409,11 +346,14 @@ func (m *Manager) allocatePrediction(runner *Runner, req PredictionRequest) { //
 		defer func() {
 			// When watcher exits, handle terminal webhook and cleanup
 			pending.mu.Lock()
-			finalResponse := pending.response
+			pending.response.populateFromRequest(pending.request)
+			if err := pending.response.finalizeResponse(); err != nil {
+				log.Errorw("failed to finalize response", "error", err)
+			}
 
 			// Send terminal webhook if prediction completed
-			if finalResponse.Status.IsCompleted() && pending.terminalWebhookSent.CompareAndSwap(false, true) {
-				_ = m.sendTerminalWebhook(pending.request, finalResponse)
+			if err := m.sendTerminalWebhook(pending); err != nil {
+				log.Errorw("failed to send terminal webhook", "error", err)
 			}
 			pending.mu.Unlock()
 
@@ -423,7 +363,7 @@ func (m *Manager) allocatePrediction(runner *Runner, req PredictionRequest) { //
 			runner.mu.Unlock()
 
 			// In one-shot mode, stop runner after prediction completes to trigger cleanup
-			if m.cfg.OneShot && finalResponse.Status.IsCompleted() {
+			if m.cfg.OneShot {
 				go func() {
 					logger := m.logger.Sugar()
 					logger.Infow("one-shot mode: stopping runner after prediction completion", "prediction_id", req.ID, "runner_id", runner.runnerCtx.id)
@@ -461,7 +401,7 @@ func (m *Manager) allocatePrediction(runner *Runner, req PredictionRequest) { //
 	}()
 }
 
-func (m *Manager) assignReqToRunner(ctx context.Context, req PredictionRequest) (*Runner, error) {
+func (m *Manager) assignReqToRunner(ctx context.Context, req PredictionRequest, async bool) (*Runner, error) {
 	log := m.logger.Sugar()
 
 	if !m.cfg.UseProcedureMode {
@@ -523,6 +463,11 @@ func (m *Manager) assignReqToRunner(ctx context.Context, req PredictionRequest)
 	m.allocatePrediction(procRunner, req) //nolint:contextcheck // see above note
 	m.mu.Unlock()
 
+	if !async {
+		if err := waitForRunnerSetup(ctx, procRunner); err != nil {
+			return nil, err
+		}
+	}
 	return procRunner, nil
 }
 
@@ -1088,26 +1033,25 @@ func (m *Manager) monitorRunnerSubprocess(ctx context.Context, runnerName string
 			pending.mu.Unlock()
 
 			failedResponse := PredictionResponse{
-				ID:      id,
 				Status:  PredictionFailed,
-				Input:   pending.request.Input,
 				Error:   "setup failed",
 				Logs:    allLogs,
 				Metrics: pending.response.Metrics,
 			}
+			failedResponse.populateFromRequest(pending.request)
 
 			pending.safeSend(failedResponse)
 			pending.safeClose()
 
 			// Update pending response with failed response for webhook
 			pending.mu.Lock()
 			pending.response = failedResponse
-			pending.mu.Unlock()
 
 			// Send terminal webhook since we're canceling the watcher
-			if pending.terminalWebhookSent.CompareAndSwap(false, true) {
-				_ = pending.sendWebhookSync(webhook.EventCompleted)
+			if err := m.sendTerminalWebhook(pending); err != nil {
+				log.Errorw("failed to send terminal webhook", "error", err)
 			}
+			pending.mu.Unlock()
 
 			for _, inputPath := range pending.inputPaths {
 				if err := os.Remove(inputPath); err != nil {
@@ -1143,26 +1087,25 @@ func (m *Manager) monitorRunnerSubprocess(ctx context.Context, runnerName string
 			pending.mu.Unlock()
 
 			failedResponse := PredictionResponse{
-				ID:      id,
 				Status:  PredictionFailed,
-				Input:   pending.request.Input,
 				Error:   "prediction failed",
 				Logs:    allLogs,
 				Metrics: pending.response.Metrics,
 			}
+			failedResponse.populateFromRequest(pending.request)
 
 			pending.safeSend(failedResponse)
 			pending.safeClose()
 
 			// Update pending response with failed response for webhook
 			pending.mu.Lock()
 			pending.response = failedResponse
-			pending.mu.Unlock()
 
 			// Send terminal webhook since we're canceling the watcher
-			if pending.terminalWebhookSent.CompareAndSwap(false, true) {
-				_ = pending.sendWebhookSync(webhook.EventCompleted)
+			if err := m.sendTerminalWebhook(pending); err != nil {
+				log.Errorw("failed to send terminal webhook", "error", err)
 			}
+			pending.mu.Unlock()
 
 			for _, inputPath := range pending.inputPaths {
 				if err := os.Remove(inputPath); err != nil {
@@ -1180,3 +1123,15 @@ func (m *Manager) monitorRunnerSubprocess(ctx context.Context, runnerName string
 		runner.status = StatusDefunct
 	}
 }
+
+func (m *Manager) sendTerminalWebhook(pending *PendingPrediction) error {
+	log := m.logger.Sugar()
+	// Send terminal webhook since we're canceling the watcher
+	if pending.response.Status.IsCompleted() && pending.terminalWebhookSent.CompareAndSwap(false, true) {
+		if err := pending.response.finalizeResponse(); err != nil {
+			log.Errorw("failed to finalize response", "error", err)
+		}
+		return pending.sendWebhookSync(webhook.EventCompleted)
+	}
+	return nil
+}

internal/runner/manager_test.go

@@ -579,7 +579,7 @@ func TestManagerPredictionHandling(t *testing.T) {
 		require.NoError(t, err)
 
 		req := PredictionRequest{ID: "test-id"}
-		_, err = m.Predict(req)
+		_, err = m.PredictSync(req)
 		require.Error(t, err)
 		assert.Equal(t, ErrNoCapacity, err)
 	})