diff --git a/sdk/api/handlers/handlers_interceptors.go b/sdk/api/handlers/handlers_interceptors.go --- a/sdk/api/handlers/handlers_interceptors.go +++ b/sdk/api/handlers/handlers_interceptors.go @@ -1,6 +1,7 @@ package handlers import ( + stdcontext "context" "net/http" "sync" "time" @@ -66,6 +67,9 @@ type requestLifecycleSkipHost interface { type requestLifecycleTracker struct { once sync.Once + cancelMu sync.Mutex + cancelStop func() bool + completed bool ctx context.Context host PluginInterceptorHost skipPluginID string @@ -76,7 +80,7 @@ func (h *BaseAPIHandler) newRequestLifecycleTracker(ctx context.Context, sourceF requestID := uuid.NewString() traceID := logging.GetRequestID(ctx) ctx = coreusage.WithRequestIdentity(ctx, requestID, traceID) - return &requestLifecycleTracker{ + tracker := &requestLifecycleTracker{ ctx: ctx, host: h.interceptorHost(), skipPluginID: skipPluginID, @@ -91,6 +95,39 @@ func (h *BaseAPIHandler) newRequestLifecycleTracker(ctx context.Context, sourceF Metadata: metadata, }, } + tracker.watchCancellation(ctx) + return tracker +} + +func (t *requestLifecycleTracker) watchCancellation(ctx context.Context) { + if t == nil || ctx == nil || ctx.Done() == nil { + return + } + stop := stdcontext.AfterFunc(ctx, func() { + t.complete(pluginapi.RequestCompletionCanceled, 0, ctx.Err()) + }) + t.cancelMu.Lock() + if t.completed { + t.cancelMu.Unlock() + stop() + return + } + t.cancelStop = stop + t.cancelMu.Unlock() +} + +func (t *requestLifecycleTracker) stopCancellationWatch() { + if t == nil { + return + } + t.cancelMu.Lock() + t.completed = true + stop := t.cancelStop + t.cancelStop = nil + t.cancelMu.Unlock() + if stop != nil { + stop() + } } func (t *requestLifecycleTracker) executionContext() context.Context { @@ -112,6 +149,7 @@ func (t *requestLifecycleTracker) complete(outcome pluginapi.RequestCompletionOu return } t.once.Do(func() { + t.stopCancellationWatch() completion := t.completion completion.Outcome = outcome completion.StatusCode = statusCode diff --git a/sdk/api/handlers/handlers_interceptors_test.go b/sdk/api/handlers/handlers_interceptors_test.go --- a/sdk/api/handlers/handlers_interceptors_test.go +++ b/sdk/api/handlers/handlers_interceptors_test.go @@ -516,6 +516,78 @@ func TestHandlerLifecycleCompletesCanceledStream(t *testing.T) { } } +func TestHandlerLifecycleCompletesCanceledStreamWhileExecutorIsStarting(t *testing.T) { + model := "handler-interceptor-lifecycle-canceled-before-stream-start" + executorStarted := make(chan struct{}) + releaseExecutor := make(chan struct{}) + executor := &interceptorCaptureExecutor{ + stream: func(context.Context, *coreauth.Auth, coreexecutor.Request, coreexecutor.Options) (*coreexecutor.StreamResult, error) { + close(executorStarted) + <-releaseExecutor + chunks := make(chan coreexecutor.StreamChunk) + close(chunks) + return &coreexecutor.StreamResult{Chunks: chunks}, nil + }, + } + handler := newInterceptorHandler(t, model, executor, &sdkconfig.SDKConfig{}) + completions := make(chan pluginapi.RequestCompletion, 2) + handler.SetPluginHost(&handlerInterceptorTestHost{ + completeRequest: func(_ context.Context, completion pluginapi.RequestCompletion) { + completions <- completion + }, + }) + + ctx, cancel := context.WithCancel(context.Background()) + type streamResult struct { + data <-chan []byte + errs <-chan *interfaces.ErrorMessage + } + result := make(chan streamResult, 1) + go func() { + data, _, errs := handler.ExecuteStreamWithAuthManager(ctx, "openai", model, []byte(`{"model":"`+model+`","stream":true}`), "") + result <- streamResult{data: data, errs: errs} + }() + + select { + case <-executorStarted: + case <-time.After(time.Second): + close(releaseExecutor) + t.Fatal("executor did not start") + } + cancel() + + select { + case completion := <-completions: + if completion.Outcome != pluginapi.RequestCompletionCanceled || completion.StatusCode != 0 || completion.Error == "" { + close(releaseExecutor) + t.Fatalf("completion = %#v", completion) + } + case <-time.After(time.Second): + close(releaseExecutor) + t.Fatal("missing cancellation completion while executor startup was blocked") + } + + close(releaseExecutor) + channels := <-result + for channels.data != nil || channels.errs != nil { + select { + case _, ok := <-channels.data: + if !ok { + channels.data = nil + } + case _, ok := <-channels.errs: + if !ok { + channels.errs = nil + } + } + } + select { + case duplicate := <-completions: + t.Fatalf("duplicate stream completion = %#v", duplicate) + default: + } +} + func TestHandlerRequestInterceptorRewritesExecutorRequest(t *testing.T) { model := "handler-interceptor-request-model" executor := &interceptorCaptureExecutor{}