Retain execution controls in prepared handles

execution_timeout_contract_test.go

@@ -38,3 +38,44 @@ func TestRunProviderDeadlineIsNotOperationTimeout(t *testing.T) {
 		t.Fatalf("provider deadline attribution: %v", err)
 	}
 }
+
+func TestPreparedExecutionRetainsTimeoutWithoutIdleTimer(t *testing.T) {
+	synctest.Test(t, func(t *testing.T) {
+		engine := observationEngine(t, observationClientFunc(func(ctx context.Context, _ promptkit.GenerateRequest) (*promptkit.GenerateResponse, error) {
+			<-ctx.Done()
+			return nil, ctx.Err()
+		}))
+		req := promptkit.RunRequest{PromptID: "prompt", OperationTimeout: time.Second}
+		handle, err := engine.PrepareExecution(context.Background(), req)
+		if err != nil {
+			t.Fatal(err)
+		}
+		req.OperationTimeout = 0
+		time.Sleep(2 * time.Second)
+		start := time.Now()
+		_, err = engine.RunPrepared(context.Background(), handle)
+		if time.Since(start) != time.Second || !errors.Is(err, promptkit.ErrOperationTimeout) || !errors.Is(err, context.DeadlineExceeded) || !errors.Is(err, promptkit.ErrLLMGenerate) {
+			t.Fatalf("elapsed=%v error=%v", time.Since(start), err)
+		}
+	})
+}
+
+func TestPreparedExecutionRespectsCallerDeadline(t *testing.T) {
+	synctest.Test(t, func(t *testing.T) {
+		engine := observationEngine(t, observationClientFunc(func(ctx context.Context, _ promptkit.GenerateRequest) (*promptkit.GenerateResponse, error) {
+			<-ctx.Done()
+			return nil, ctx.Err()
+		}))
+		ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
+		defer cancel()
+		handle, err := engine.PrepareExecution(ctx, promptkit.RunRequest{PromptID: "prompt", OperationTimeout: 5 * time.Second})
+		if err != nil {
+			t.Fatal(err)
+		}
+		time.Sleep(time.Second)
+		_, err = engine.RunPrepared(ctx, handle)
+		if !errors.Is(err, context.DeadlineExceeded) || errors.Is(err, promptkit.ErrOperationTimeout) {
+			t.Fatalf("caller deadline attribution: %v", err)
+		}
+	})
+}