fix(sdk): include externalized reads in the intercept budget

sdk/go/intercept.go

@@ -0,0 +1,58 @@
+package extension
+
+import (
+	"context"
+	"encoding/json"
+	"errors"
+	"fmt"
+	"time"
+)
+
+func (s *server) handleIntercept(ctx context.Context, raw json.RawMessage) (any, error) {
+	var p InterceptParams
+	if err := strictDecode(raw, &p); err != nil {
+		return nil, MustProtocolError(ErrInvalidParams)
+	}
+	if !validInterceptEvent(p.Event) || p.Seq < 1 || p.TimeoutMillis < 0 || !jsonKeyPresent(raw, "payload") {
+		return nil, MustProtocolError(ErrInvalidParams)
+	}
+	if p.TimeoutMillis > 0 {
+		var cancel context.CancelFunc
+		ctx, cancel = context.WithTimeout(ctx, time.Duration(p.TimeoutMillis)*time.Millisecond)
+		defer cancel()
+	}
+	payload, err := s.rehydrate(ctx, p.Payload, p.Externalized, "/payload")
+	if err == nil {
+		err = ctx.Err()
+	}
+	if err != nil {
+		if errors.Is(err, context.DeadlineExceeded) && errors.Is(ctx.Err(), context.DeadlineExceeded) {
+			return nil, MustProtocolError(ErrInterceptTimeout)
+		}
+		return nil, err
+	}
+	fn := s.opts.Interceptors[string(p.Event)]
+	if fn == nil {
+		fn = s.opts.Interceptors["*"]
+	}
+	if fn == nil {
+		return Continue(), nil
+	}
+	result, err := fn(ctx, string(p.Event), payload)
+	if err != nil {
+		// The callback's advertised intercept budget expired. Return the
+		// frozen timeout reason rather than racing the host's identical timer
+		// with a generic internal error response.
+		if errors.Is(err, context.DeadlineExceeded) && errors.Is(ctx.Err(), context.DeadlineExceeded) {
+			return nil, MustProtocolError(ErrInterceptTimeout)
+		}
+		return nil, err
+	}
+	if result == nil {
+		return Continue(), nil
+	}
+	if !validInterceptDecision(result.Decision) {
+		return nil, fmt.Errorf("extension: interceptor for %q returned invalid decision %q", p.Event, result.Decision)
+	}
+	return result, nil
+}

sdk/go/intercept_deadline_test.go

@@ -0,0 +1,69 @@
+package extension
+
+import (
+	"context"
+	"encoding/json"
+	"strings"
+	"sync/atomic"
+	"testing"
+	"time"
+)
+
+func TestInterceptBudgetIncludesExternalizedContent(t *testing.T) {
+	for _, tc := range []struct {
+		name    string
+		timeout int
+		failed  bool
+	}{
+		{"bounded", 20, true},
+		{"unbounded", 0, false},
+		{"shared budget", 1000, false},
+	} {
+		t.Run(tc.name, func(t *testing.T) {
+			store := newContentStore()
+			payload := `{"text":"` + strings.Repeat("x", ExternalizeFieldBytes+1) + `"}`
+			_, descriptor := store.put(payload)
+			var called atomic.Bool
+			var firstRead atomic.Pointer[time.Time]
+			var got string
+			host, _ := startFakeHost(t, basicHandler(), Options{Interceptors: map[string]InterceptorFunc{
+				"input.receive": func(ctx context.Context, _ string, raw json.RawMessage) (*InterceptResult, error) {
+					if tc.timeout > 0 {
+						deadline, ok := ctx.Deadline()
+						started := firstRead.Load()
+						if !ok || started == nil || deadline.After(started.Add(time.Duration(tc.timeout)*time.Millisecond)) {
+							t.Error("callback received a new budget after content reading")
+						}
+					}
+					got = string(raw)
+					called.Store(true)
+					return Continue(), nil
+				},
+			}})
+			host.onRequest(MethodHostContentRead, func(params json.RawMessage) (any, *hostError) {
+				started := time.Now()
+				firstRead.CompareAndSwap(nil, &started)
+				time.Sleep(100 * time.Millisecond)
+				return store.handler(params)
+			})
+			host.handshake(t)
+			resp := host.request(MethodExtensionIntercept, InterceptParams{
+				Event: EventInputReceive, Seq: 1, Payload: json.RawMessage("null"), TimeoutMillis: tc.timeout,
+				Externalized: []ExternalizedField{descriptor},
+			})
+			if tc.failed {
+				if resp.Err == nil || resp.Err.Code != DomainErrorCode {
+					t.Fatalf("expected content-read timeout, got %+v", resp)
+				}
+				data, _ := resp.Err.Data.(ProtocolErrorData)
+				if data.Reason != ErrInterceptTimeout || called.Load() {
+					t.Fatalf("reason = %q, callback called = %v", data.Reason, called.Load())
+				}
+				return
+			}
+			if resp.Err != nil || !called.Load() || got != payload {
+				t.Fatalf("unbounded intercept = %+v, called = %v, bytes = %d", resp, called.Load(), len(got))
+			}
+		})
+	}
+}

sdk/go/sdk.go

@@ -490,49 +490,6 @@ func (s *server) handleShutdown(ctx context.Context, raw json.RawMessage) (any,
 
 // Intercept and observation
 
-func (s *server) handleIntercept(ctx context.Context, raw json.RawMessage) (any, error) {
-	var p InterceptParams
-	if err := strictDecode(raw, &p); err != nil {
-		return nil, MustProtocolError(ErrInvalidParams)
-	}
-	if !validInterceptEvent(p.Event) || p.Seq < 1 || p.TimeoutMillis < 0 || !jsonKeyPresent(raw, "payload") {
-		return nil, MustProtocolError(ErrInvalidParams)
-	}
-	payload, err := s.rehydrate(ctx, p.Payload, p.Externalized, "/payload")
-	if err != nil {
-		return nil, err
-	}
-	fn := s.opts.Interceptors[string(p.Event)]
-	if fn == nil {
-		fn = s.opts.Interceptors["*"]
-	}
-	if fn == nil {
-		return Continue(), nil
-	}
-	if p.TimeoutMillis > 0 {
-		var cancel context.CancelFunc
-		ctx, cancel = context.WithTimeout(ctx, time.Duration(p.TimeoutMillis)*time.Millisecond)
-		defer cancel()
-	}
-	result, err := fn(ctx, string(p.Event), payload)
-	if err != nil {
-		// The callback's advertised intercept budget expired. Return the
-		// frozen timeout reason rather than racing the host's identical timer
-		// with a generic internal error response.
-		if errors.Is(err, context.DeadlineExceeded) && errors.Is(ctx.Err(), context.DeadlineExceeded) {
-			return nil, MustProtocolError(ErrInterceptTimeout)
-		}
-		return nil, err
-	}
-	if result == nil {
-		return Continue(), nil
-	}
-	if !validInterceptDecision(result.Decision) {
-		return nil, fmt.Errorf("extension: interceptor for %q returned invalid decision %q", p.Event, result.Decision)
-	}
-	return result, nil
-}
-
 func (s *server) handleEvent(ctx context.Context, raw json.RawMessage) {
 	var p EventParams
 	if err := strictDecode(raw, &p); err != nil || !validInterceptEvent(p.Event) || !jsonKeyPresent(raw, "payload") {